60 lines
2.5 KiB
Python
60 lines
2.5 KiB
Python
"""Unit tests for LLM client and Repository — with fakes."""
|
|
from services.llm_client import FakeLLMClient
|
|
from repository import MemRepository
|
|
|
|
|
|
class TestFakeLLM:
|
|
def test_exact_match(self):
|
|
client = FakeLLMClient({"hello": "world"})
|
|
assert client.complete("hello") == "world"
|
|
|
|
def test_partial_match(self):
|
|
client = FakeLLMClient({"classify": '{"doc_type":"contract"}'})
|
|
assert "contract" in client.complete("classify this document")
|
|
|
|
def test_default_fallback(self):
|
|
client = FakeLLMClient({"default": '{"ok":true}'})
|
|
assert client.complete("unknown prompt") == '{"ok":true}'
|
|
|
|
def test_call_history(self):
|
|
client = FakeLLMClient({"a": "1", "b": "2"})
|
|
client.complete("a")
|
|
client.complete("b")
|
|
assert len(client.calls) == 2
|
|
assert client.calls[0] == "a"
|
|
|
|
|
|
class TestMemRepository:
|
|
def test_insert_and_retrieve(self, mem_repo):
|
|
doc = mem_repo.insert_document("test.docx", "application/vnd...", "base64data", batch_id="batch-1")
|
|
assert doc["filename"] == "test.docx"
|
|
assert doc["classify_status"] == "pending"
|
|
|
|
def test_list_pending(self, mem_repo):
|
|
mem_repo.insert_document("a.docx", "mime", "data", batch_id="b1")
|
|
mem_repo.insert_document("b.docx", "mime", "data", batch_id="b1")
|
|
mem_repo.insert_document("c.docx", "mime", "data", batch_id="b2")
|
|
pending = mem_repo.list_pending("b1")
|
|
assert len(pending) == 2
|
|
|
|
def test_set_classification(self, mem_repo):
|
|
doc = mem_repo.insert_document("test.docx", "mime", "data", batch_id="b1")
|
|
mem_repo.set_classification(doc["id"], "contract", "03700", None, "2026-02-01", "XXX")
|
|
updated = mem_repo.documents[doc["id"]]
|
|
assert updated["doc_type"] == "contract"
|
|
assert updated["classify_status"] == "classified"
|
|
|
|
def test_set_garbage(self, mem_repo):
|
|
doc = mem_repo.insert_document("счет.docx", "mime", "data", batch_id="b1")
|
|
mem_repo.set_classify_garbage(doc["id"], "filename_regex")
|
|
assert mem_repo.documents[doc["id"]]["doc_type"] == "garbage"
|
|
|
|
def test_contract_lifecycle(self, mem_repo):
|
|
cid = mem_repo.insert_contract("03700_1", "ЗАО XXX")
|
|
assert cid
|
|
doc = mem_repo.insert_document("test.docx", "mime", "data")
|
|
mem_repo.insert_supplement(cid, doc["id"], "initial")
|
|
supps = mem_repo.list_supplements(cid)
|
|
assert len(supps) == 1
|
|
assert supps[0]["type"] == "initial"
|