57 lines
2.5 KiB
Python
57 lines
2.5 KiB
Python
|
"""Unit tests for the startup model invariant -- no real DB."""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import pathlib
|
||
|
|
import sys
|
||
|
|
|
||
|
|
import pytest
|
||
|
|
|
||
|
|
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1]))
|
||
|
|
|
||
|
|
from app.startup import ModelInvariantError, validate_embed_model # noqa: E402
|
||
|
|
|
||
|
|
|
||
|
|
class _FakeConn:
|
||
|
|
def __init__(self, chunk_models: set, summary_embedding_models: set):
|
||
|
|
self._chunk_models = chunk_models
|
||
|
|
self._summary_embedding_models = summary_embedding_models
|
||
|
|
|
||
|
|
async def fetch(self, query, *params):
|
||
|
|
if "FROM document_chunk" in query:
|
||
|
|
return [{"model": m} for m in self._chunk_models]
|
||
|
|
if "FROM document_summary" in query:
|
||
|
|
return [{"embedding_model": m} for m in self._summary_embedding_models]
|
||
|
|
raise AssertionError(f"unexpected query: {query}")
|
||
|
|
|
||
|
|
|
||
|
|
class TestValidateEmbedModel:
|
||
|
|
async def test_passes_when_model_present_in_both_tables(self):
|
||
|
|
conn = _FakeConn(chunk_models={"bge-m3"}, summary_embedding_models={"bge-m3"})
|
||
|
|
await validate_embed_model(conn, "bge-m3") # must not raise
|
||
|
|
|
||
|
|
async def test_fails_when_chunk_model_missing(self):
|
||
|
|
conn = _FakeConn(chunk_models={"some-other-model"}, summary_embedding_models={"bge-m3"})
|
||
|
|
with pytest.raises(ModelInvariantError, match="document_chunk"):
|
||
|
|
await validate_embed_model(conn, "bge-m3")
|
||
|
|
|
||
|
|
async def test_fails_when_chunk_table_empty(self):
|
||
|
|
conn = _FakeConn(chunk_models=set(), summary_embedding_models={"bge-m3"})
|
||
|
|
with pytest.raises(ModelInvariantError, match="document_chunk"):
|
||
|
|
await validate_embed_model(conn, "bge-m3")
|
||
|
|
|
||
|
|
async def test_fails_when_summary_embedding_model_missing(self):
|
||
|
|
conn = _FakeConn(chunk_models={"bge-m3"}, summary_embedding_models={"some-other-model"})
|
||
|
|
with pytest.raises(ModelInvariantError, match="document_summary"):
|
||
|
|
await validate_embed_model(conn, "bge-m3")
|
||
|
|
|
||
|
|
async def test_does_not_confuse_summary_writer_model_with_embedding_model(self):
|
||
|
|
# Regression guard: document_summary.model is the LLM that WROTE the summary
|
||
|
|
# (claude-haiku-4-5/gemma3:12b), never bge-m3 -- checking that column instead of
|
||
|
|
# embedding_model would make this invariant impossible to satisfy.
|
||
|
|
conn = _FakeConn(
|
||
|
|
chunk_models={"bge-m3"},
|
||
|
|
summary_embedding_models={"bge-m3"},
|
||
|
|
)
|
||
|
|
assert conn._summary_embedding_models == {"bge-m3"}
|
||
|
|
await validate_embed_model(conn, "bge-m3") # must not raise
|