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

191 lines
7.4 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-1 heuristic extractor (regex-free port).
Ports 100% of the regex-v1 fixtures from test_prescreen_extractor.py onto
HeuristicExtractor (plan Phase 1) and adds coverage for the string helpers
(`_scan_number`, header window, trigger scans).
"""
from __future__ import annotations
import pytest
from contract_check.worker_prescreen.extractor_heuristic import (
EXTRACTOR_VERSION,
HeuristicExtractor,
_header_window,
_scan_number,
)
SIMPLE_SUPPLY = """\
ДОГОВОР ПОСТАВКИ № 42
г. Москва
Общество с ограниченной ответственностью «Продавец»,
именуемое в дальнейшем «Поставщик», в лице директора Иванова И.И.,
с одной стороны, и
Общество с ограниченной ответственностью «Покупатель»,
именуемое в дальнейшем «Покупатель», с другой стороны,
заключили настоящий договор о нижеследующем.
1. Предмет договора
Поставщик обязуется передать в собственность Покупателю товар.
2. Цена и порядок расчётов
2.1. Общая стоимость товара составляет 1 250 000 (один миллион двести пятьдесят тысяч) рублей.
3. Срок действия договора
3.1. Настоящий договор вступает в силу с 01.09.2025 и действует по 31.08.2026.
4. Ответственность сторон
4.1. За просрочку поставки Поставщик уплачивает неустойку в размере 0,1% от стоимости.
5. Порядок разрешения споров
5.1. Споры подлежат рассмотрению в Арбитражном суде г. Москвы.
"""
MINIMAL = """\
Договор оказания услуг между ИП Петров Петр Петрович и ООО «Клиент».
Стоимость услуг составляет 50 000 рублей.
"""
@pytest.fixture
def extractor() -> HeuristicExtractor:
return HeuristicExtractor()
def test_version_is_heuristic_v2() -> None:
assert EXTRACTOR_VERSION == "heuristic-v2"
assert HeuristicExtractor.extractor_version == "heuristic-v2"
def test_supply_extraction(extractor: HeuristicExtractor) -> None:
meta = extractor.extract(SIMPLE_SUPPLY)
assert meta.contract_type == "supply"
assert "Продавец" in (meta.party_a or "")
assert "Покупатель" in (meta.party_b or "")
assert meta.total_amount == 1250000.0
assert meta.currency == "RUB"
assert meta.start_date == "2025-09-01"
assert meta.end_date == "2026-08-31"
assert meta.has_penalty_clause is True
assert meta.has_termination_clause is False
assert meta.has_arbitration is True
assert meta.confidence_score > 0.6
def test_minimal_extraction(extractor: HeuristicExtractor) -> None:
meta = extractor.extract(MINIMAL)
assert meta.contract_type == "services"
assert "ИП Петров" in (meta.party_a or "")
assert "ООО «Клиент»" in (meta.party_b or "")
assert meta.total_amount == 50000.0
assert meta.currency == "RUB"
def test_confidence_score_bounds(extractor: HeuristicExtractor) -> None:
meta = extractor.extract("")
assert meta.confidence_score == 0.0
assert 0.0 <= meta.confidence_score <= 1.0
def test_party_extraction_handles_whitespace(extractor: HeuristicExtractor) -> None:
text = "Договор между ООО «Альфа» и ООО «Бета»"
meta = extractor.extract(text)
assert "Альфа" in (meta.party_a or "")
assert "Бета" in (meta.party_b or "")
@pytest.mark.parametrize(
("clause", "expected"),
[
("Стороны вправе расторгнуть договор в одностороннем порядке", "has_termination_clause"),
("В случае нарушения Покупатель уплачивает штраф", "has_penalty_clause"),
("Споры рассматриваются в Арбитражном суде", "has_arbitration"),
],
)
def test_boolean_flags(clause: str, expected: str, extractor: HeuristicExtractor) -> None:
meta = extractor.extract(f"Договор поставки. {clause}.")
assert getattr(meta, expected) is True
# ── helper coverage ──────────────────────────────────────────────────────────
@pytest.mark.parametrize(
("raw", "expected"),
[
("1 250 000 рублей", 1250000.0),
("1250000", 1250000.0),
("50 000 руб", 50000.0),
("0,1% от стоимости", 0.1),
("1 250 000,50", 1250000.5),
("999", 999.0),
],
)
def test_scan_number(raw: str, expected: float) -> None:
parsed = _scan_number(raw, 0)
assert parsed is not None
assert parsed[0] == expected
def test_scan_number_rejects_non_digits() -> None:
assert _scan_number("abc", 0) is None
def test_scan_number_nbsp_thousands() -> None:
parsed = _scan_number("1\u00a0250\u00a0000", 0)
assert parsed is not None
assert parsed[0] == 1250000.0
def test_header_window_cuts_at_numbered_heading() -> None:
text = "ДОГОВОР № 1\nООО «Альфа» и ООО «Бета»\n1. Предмет договора\nООО «Гамма»"
window = _header_window(text)
assert "Альфа" in window
assert "Бета" in window
assert "Гамма" not in window
def test_parties_only_scanned_in_header(extractor: HeuristicExtractor) -> None:
text = (
"ДОГОВОР № 1\nООО «Альфа» и ООО «Бета»\n\n"
"1. Предмет договора\nПоставщик обязуется...\n\n"
"10. Реквизиты сторон\nООО «Гамма»"
)
meta = extractor.extract(text)
assert "Альфа" in (meta.party_a or "")
assert "Бета" in (meta.party_b or "")
def test_end_date_only_via_do_fallback(extractor: HeuristicExtractor) -> None:
meta = extractor.extract("Соглашение действует до 31.12.2026.")
assert meta.start_date is None
assert meta.end_date == "2026-12-31"
def test_start_date_via_ot_token(extractor: HeuristicExtractor) -> None:
meta = extractor.extract("Договор заключён от 15.03.2024.")
assert meta.start_date == "2024-03-15"
def test_amount_skips_trigger_without_digits(extractor: HeuristicExtractor) -> None:
meta = extractor.extract("Общая стоимость товара определяется расчётом. Сумма 300 BYN.")
assert meta.total_amount == 300.0
assert meta.currency == "BYN"
def test_no_false_entity_match_inside_words(extractor: HeuristicExtractor) -> None:
meta = extractor.extract("Ипотека оформлена. Сумма 100 рублей.")
assert meta.party_a is None
assert meta.party_b is None
def test_bare_entity_form_without_name_dropped(extractor: HeuristicExtractor) -> None:
meta = extractor.extract("Договор между ООО и АО.\nСумма 100 рублей.")
assert meta.party_a is None
assert meta.party_b is None