microdao-daarion/tests/test_stepan_doc_anchor_reset.py

"""Tests for Doc Anchor Reset (PROMPT 27, v3.3)."""
import sys
import os
import time

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from crews.agromatrix_crew import session_context as sc
from crews.agromatrix_crew.doc_facts import build_self_correction

FACTS_DOC_A = {
    "profit_uah": 5_972_016.0,
    "fertilizer_uah": 1_521_084.0,
    "area_ha": 497.0,
}


def _make_session(active_doc_id: str, doc_facts: dict, fact_claims: list | None = None) -> dict:
    sess = sc._default_session()
    sess["active_doc_id"] = active_doc_id
    sess["doc_facts"] = dict(doc_facts)
    sess["fact_claims"] = list(fact_claims or [])
    sess["updated_at"] = time.time()
    return sess


def _store_session(chat_id: str, sess: dict):
    sc._STORE[chat_id] = sess


# ── Основний сценарій: зміна doc_id → reset ──────────────────────────────────

def test_anchor_reset_clears_doc_facts():
    """При зміні doc_id — doc_facts скидаються."""
    chat_id = "anchor_test_1"
    sess = _make_session("docA", FACTS_DOC_A)
    _store_session(chat_id, sess)

    loaded = sc.load_session(chat_id)
    assert loaded["doc_facts"] == FACTS_DOC_A

    # Симулюємо логіку з run.py: новий doc_id != старий
    current_doc_id = "docB"
    prev_doc_id = loaded.get("active_doc_id")
    if prev_doc_id and prev_doc_id != current_doc_id:
        loaded["doc_facts"] = {}
        loaded["fact_claims"] = []

    assert loaded["doc_facts"] == {}
    assert loaded["fact_claims"] == []


def test_anchor_reset_clears_fact_claims():
    """При зміні doc_id — fact_claims також скидаються."""
    chat_id = "anchor_test_2"
    claims = [{"key": "profit_present", "value": True, "ts": time.time()}]
    sess = _make_session("docA", FACTS_DOC_A, fact_claims=claims)
    _store_session(chat_id, sess)

    loaded = sc.load_session(chat_id)
    assert len(loaded["fact_claims"]) == 1

    # Зміна doc_id
    current_doc_id = "docB"
    prev_doc_id = loaded.get("active_doc_id")
    if prev_doc_id and prev_doc_id != current_doc_id:
        loaded["doc_facts"] = {}
        loaded["fact_claims"] = []

    assert loaded["fact_claims"] == []


def test_anchor_reset_updates_active_doc_id():
    """Після reset active_doc_id має оновитися через update_session."""
    chat_id = "anchor_test_3"
    sess = _make_session("docA", FACTS_DOC_A)
    _store_session(chat_id, sess)

    # Виконуємо update_session з новим active_doc_id
    sc.update_session(chat_id, "текст", depth="deep", agents=[],
                      active_doc_id="docB", doc_facts={}, fact_claims=[])

    updated = sc.load_session(chat_id)
    assert updated["active_doc_id"] == "docB"
    assert updated["doc_facts"] == {}


# ── Fail-safe: якщо current_doc_id None → нічого не скидаємо ────────────────

def test_no_reset_when_current_doc_id_is_none():
    """Якщо current_doc_id не визначений — doc_facts НЕ скидаються."""
    chat_id = "anchor_test_failsafe"
    sess = _make_session("docA", FACTS_DOC_A)
    _store_session(chat_id, sess)

    loaded = sc.load_session(chat_id)
    current_doc_id = None  # невідомий

    # Логіка з run.py: якщо current_doc_id None — нічого не робимо
    if current_doc_id:
        prev = loaded.get("active_doc_id")
        if prev and prev != current_doc_id:
            loaded["doc_facts"] = {}
            loaded["fact_claims"] = []

    # doc_facts мають залишитись
    assert loaded["doc_facts"] == FACTS_DOC_A


def test_no_reset_when_same_doc_id():
    """Якщо doc_id той самий — doc_facts НЕ скидаються."""
    chat_id = "anchor_test_same"
    sess = _make_session("docA", FACTS_DOC_A)
    _store_session(chat_id, sess)

    loaded = sc.load_session(chat_id)
    current_doc_id = "docA"  # той самий

    prev = loaded.get("active_doc_id")
    if prev and prev != current_doc_id:
        loaded["doc_facts"] = {}
        loaded["fact_claims"] = []

    assert loaded["doc_facts"] == FACTS_DOC_A


def test_no_reset_when_no_previous_doc_id():
    """Якщо попереднього active_doc_id нема — reset не спрацьовує."""
    chat_id = "anchor_test_noprev"
    sess = sc._default_session()
    sess["doc_facts"] = FACTS_DOC_A
    sess["active_doc_id"] = None  # нема попереднього
    sess["updated_at"] = time.time()
    _store_session(chat_id, sess)

    loaded = sc.load_session(chat_id)
    current_doc_id = "docB"

    prev = loaded.get("active_doc_id")
    if prev and prev != current_doc_id:
        loaded["doc_facts"] = {}
        loaded["fact_claims"] = []

    # Нема попереднього — нічого не скинули
    assert loaded["doc_facts"] == FACTS_DOC_A


# ── Self-correction: тільки в межах одного doc_id ────────────────────────────

def test_self_correction_blocked_on_different_doc():
    """Self-correction НЕ спрацьовує при зміні doc_id."""
    sess = _make_session("docA", FACTS_DOC_A)
    sess["fact_claims"] = [{"key": "profit_present", "value": False, "ts": time.time()}]

    response = "У звіті є прибуток — 5 972 016 грн."
    # current_doc_id = docB (новий) → correction заблокована
    correction = build_self_correction(response, FACTS_DOC_A, sess, current_doc_id="docB")
    assert correction == "", f"Expected no correction, got: {correction!r}"


def test_self_correction_works_same_doc():
    """Self-correction спрацьовує в межах того самого doc_id."""
    sess = _make_session("docA", FACTS_DOC_A)
    sess["fact_claims"] = [{"key": "profit_present", "value": False, "ts": time.time()}]

    response = "У звіті є прибуток — 5 972 016 грн."
    correction = build_self_correction(response, FACTS_DOC_A, sess, current_doc_id="docA")
    assert correction, f"Expected correction prefix, got empty"
    assert "неточно" in correction.lower() or "написав" in correction.lower()


def test_self_correction_works_when_no_current_doc_id():
    """Self-correction не блокується якщо current_doc_id=None (backward-compat)."""
    sess = _make_session("docA", FACTS_DOC_A)
    sess["fact_claims"] = [{"key": "profit_present", "value": False, "ts": time.time()}]

    response = "У звіті є прибуток — 5 972 016 грн."
    # current_doc_id=None → anchor guard не спрацьовує → correction дозволена
    correction = build_self_correction(response, FACTS_DOC_A, sess, current_doc_id=None)
    assert correction, "Expected correction when current_doc_id is None"


# ── update_session: active_doc_id персистується ───────────────────────────────

def test_update_session_persists_active_doc_id():
    """update_session зберігає active_doc_id."""
    chat_id = "anchor_persist_test"
    sc.clear_session(chat_id)

    sc.update_session(chat_id, "запит", depth="deep", agents=[],
                      active_doc_id="doc_xyz123")

    sess = sc.load_session(chat_id)
    assert sess["active_doc_id"] == "doc_xyz123"


def test_update_session_preserves_doc_id_if_not_provided():
    """update_session зберігає існуючий active_doc_id якщо новий не передано."""
    chat_id = "anchor_preserve_test"
    sc.clear_session(chat_id)

    sc.update_session(chat_id, "перший", depth="deep", agents=[],
                      active_doc_id="docA")
    sc.update_session(chat_id, "другий", depth="light", agents=[])

    sess = sc.load_session(chat_id)
    # active_doc_id має залишитись "docA" (None не перезаписує)
    assert sess["active_doc_id"] == "docA"