DealDocumentScreening/tests/unit/test_prescreen_extractor_llm.py
2026-08-17 20:49:29 +03:00

113 lines
3.3 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

"""Unit tests for the Stage-2 worker-side wrapper (dict → validated meta)."""
from __future__ import annotations
from typing import Any
from contract_check.worker_prescreen.extractor_llm import (
EXTRACTOR_VERSION,
LLMPrescreenExtractor,
validate_raw,
)
VALID: dict[str, Any] = {
"contract_type": "supply",
"party_a": "ООО «Альфа»",
"party_b": "ООО «Бета»",
"total_amount": 1250000,
"currency": "RUB",
"start_date": "2025-09-01",
"end_date": "2026-08-31",
"has_penalty_clause": True,
"has_termination_clause": False,
"has_arbitration": True,
}
def test_valid_dict_passes_through() -> None:
meta = validate_raw(VALID)
assert meta.contract_type == "supply"
assert meta.party_a == "ООО «Альфа»"
assert meta.total_amount == 1250000.0
assert meta.currency == "RUB"
assert meta.start_date == "2025-09-01"
assert meta.has_penalty_clause is True
assert meta.has_arbitration is True
assert meta.confidence_score == 1.0
def test_invalid_enum_becomes_none() -> None:
meta = validate_raw({**VALID, "contract_type": "lease-agreement"})
assert meta.contract_type is None
def test_contract_type_normalized_to_lower() -> None:
meta = validate_raw({**VALID, "contract_type": "LEASE"})
assert meta.contract_type == "lease"
def test_invalid_currency_becomes_none() -> None:
meta = validate_raw({**VALID, "currency": "CHF"})
assert meta.currency is None
def test_currency_normalized_to_upper() -> None:
meta = validate_raw({**VALID, "currency": "byn"})
assert meta.currency == "BYN"
def test_invalid_iso_date_becomes_none() -> None:
meta = validate_raw({**VALID, "start_date": "сентябрь 2025"})
assert meta.start_date is None
def test_ddmmyyyy_date_accepted() -> None:
meta = validate_raw({**VALID, "start_date": "01.09.2025"})
assert meta.start_date == "2025-09-01"
def test_impossible_date_becomes_none() -> None:
meta = validate_raw({**VALID, "start_date": "32.13.2025"})
assert meta.start_date is None
def test_negative_amount_becomes_none() -> None:
meta = validate_raw({**VALID, "total_amount": -5})
assert meta.total_amount is None
def test_non_bool_flags_become_none() -> None:
meta = validate_raw({**VALID, "has_penalty_clause": "да"})
assert meta.has_penalty_clause is None
def test_empty_strings_become_none() -> None:
meta = validate_raw({**VALID, "party_a": " "})
assert meta.party_a is None
def test_extract_version_constant() -> None:
assert EXTRACTOR_VERSION == "hybrid-llm-v1"
async def test_extractor_calls_provider_and_scores() -> None:
class StubProvider:
def __init__(self) -> None:
self.received: list[str] = []
async def analyze(self, text: str, *, checklist: str, extra_context: str = "") -> object:
raise NotImplementedError
async def extract_prescreen(self, text: str) -> dict[str, Any]:
self.received.append(text)
return dict(VALID)
async def aclose(self) -> None:
pass
provider = StubProvider()
extractor = LLMPrescreenExtractor(provider)
meta = await extractor.extract("Договор поставки...")
assert provider.received == ["Договор поставки..."]
assert meta.contract_type == "supply"
assert meta.confidence_score == 1.0