microdao-daarion/tests/test_stepan_doc_handoff.py

"""
Tests for Doc Handoff (PROMPT 28, v3.3) — chat-scoped doc_context.

Перевіряємо:
1. save_chat_doc_context / get_chat_doc_context logic (unit, з mock memory_client)
2. doc_context fallback у run.py (симуляція через session_context)
3. Інтеграційний сценарій: upload doc → питання (інший session) → Stepan бачить doc_id
"""
import sys
import os
import json
import time
import asyncio
from unittest.mock import AsyncMock, MagicMock, patch

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

# ── helpers ──────────────────────────────────────────────────────────────────

DOC_A = {
    "doc_id": "parsed_doc_abc123",
    "file_unique_id": "tg_unique_xyz",
    "file_name": "Звіт кукурудза.xlsx",
    "extracted_summary": "Площа: 497 га. Прибуток: 5 972 016 грн. Добрива: 1 521 084 грн.",
    "source": "telegram",
}


def run_async(coro):
    """Запустити корутину синхронно для тестів (Python 3.10+ safe)."""
    return asyncio.run(coro)


# ── Unit tests: save_chat_doc_context ────────────────────────────────────────

def test_save_chat_doc_context_calls_upsert():
    """save_chat_doc_context має викликати memory_client.upsert_fact з правильним ключем."""
    mock_client = AsyncMock()
    mock_client.upsert_fact = AsyncMock(return_value=True)

    with patch("services.doc_service.memory_client", mock_client):
        from services.doc_service import save_chat_doc_context
        result = run_async(save_chat_doc_context("chat_001", "agromatrix", DOC_A))

    assert result is True
    mock_client.upsert_fact.assert_called_once()
    call_kwargs = mock_client.upsert_fact.call_args.kwargs
    assert call_kwargs["user_id"] == "chat:agromatrix:chat_001"
    assert call_kwargs["fact_key"] == "doc_context_chat:agromatrix:chat_001"
    saved = call_kwargs["fact_value_json"]
    assert saved["doc_id"] == "parsed_doc_abc123"
    assert saved["file_name"] == "Звіт кукурудза.xlsx"


def test_save_chat_doc_context_truncates_summary():
    """extracted_summary обрізається до 4000 символів."""
    mock_client = AsyncMock()
    mock_client.upsert_fact = AsyncMock(return_value=True)

    big_doc = dict(DOC_A)
    big_doc["extracted_summary"] = "x" * 9000

    with patch("services.doc_service.memory_client", mock_client):
        from services.doc_service import save_chat_doc_context
        run_async(save_chat_doc_context("chat_002", "agromatrix", big_doc))

    saved = mock_client.upsert_fact.call_args.kwargs["fact_value_json"]
    assert len(saved["extracted_summary"]) == 4000


def test_save_chat_doc_context_fail_safe():
    """При помилці memory_client — повертаємо False без exception."""
    mock_client = AsyncMock()
    mock_client.upsert_fact = AsyncMock(side_effect=RuntimeError("memory down"))

    with patch("services.doc_service.memory_client", mock_client):
        from services.doc_service import save_chat_doc_context
        result = run_async(save_chat_doc_context("chat_003", "agromatrix", DOC_A))

    assert result is False


# ── Unit tests: get_chat_doc_context ─────────────────────────────────────────

def test_get_chat_doc_context_returns_dict():
    """get_chat_doc_context повертає dict якщо fact є."""
    mock_client = AsyncMock()
    mock_client.get_fact = AsyncMock(return_value={
        "fact_value_json": DOC_A,
    })

    with patch("services.doc_service.memory_client", mock_client):
        from services.doc_service import get_chat_doc_context
        result = run_async(get_chat_doc_context("chat_010", "agromatrix"))

    assert result is not None
    assert result["doc_id"] == "parsed_doc_abc123"
    assert result["file_name"] == "Звіт кукурудза.xlsx"

    call_kwargs = mock_client.get_fact.call_args.kwargs
    assert call_kwargs["user_id"] == "chat:agromatrix:chat_010"
    assert call_kwargs["fact_key"] == "doc_context_chat:agromatrix:chat_010"


def test_get_chat_doc_context_json_string_deserialized():
    """fact_value_json може прийти як JSON-рядок — deserialize правильно."""
    mock_client = AsyncMock()
    mock_client.get_fact = AsyncMock(return_value={
        "fact_value_json": json.dumps(DOC_A),
    })

    with patch("services.doc_service.memory_client", mock_client):
        from services.doc_service import get_chat_doc_context
        result = run_async(get_chat_doc_context("chat_011", "agromatrix"))

    assert result is not None
    assert result["doc_id"] == "parsed_doc_abc123"


def test_get_chat_doc_context_returns_none_when_missing():
    """Якщо fact немає — повертаємо None."""
    mock_client = AsyncMock()
    mock_client.get_fact = AsyncMock(return_value=None)

    with patch("services.doc_service.memory_client", mock_client):
        from services.doc_service import get_chat_doc_context
        result = run_async(get_chat_doc_context("chat_012", "agromatrix"))

    assert result is None


def test_get_chat_doc_context_fail_safe():
    """При помилці memory_client — повертаємо None без exception."""
    mock_client = AsyncMock()
    mock_client.get_fact = AsyncMock(side_effect=RuntimeError("timeout"))

    with patch("services.doc_service.memory_client", mock_client):
        from services.doc_service import get_chat_doc_context
        result = run_async(get_chat_doc_context("chat_013", "agromatrix"))

    assert result is None


def test_get_chat_doc_context_returns_none_without_doc_id():
    """Якщо fact є але без doc_id і file_unique_id — повертаємо None."""
    mock_client = AsyncMock()
    mock_client.get_fact = AsyncMock(return_value={
        "fact_value_json": {"file_name": "test.xlsx"},  # немає doc_id
    })

    with patch("services.doc_service.memory_client", mock_client):
        from services.doc_service import get_chat_doc_context
        result = run_async(get_chat_doc_context("chat_014", "agromatrix"))

    assert result is None


# ── Інтеграційний сценарій ───────────────────────────────────────────────────

async def _scenario_upload_then_question():
    stored = {}

    async def mock_upsert(user_id, fact_key, fact_value_json, team_id=None, **kwargs):
        stored[fact_key] = fact_value_json
        return True

    async def mock_get(user_id, fact_key, **kwargs):
        val = stored.get(fact_key)
        return {"fact_value_json": val} if val else None

    mock_client = AsyncMock()
    mock_client.upsert_fact = AsyncMock(side_effect=mock_upsert)
    mock_client.get_fact = AsyncMock(side_effect=mock_get)

    with patch("services.doc_service.memory_client", mock_client):
        from services.doc_service import save_chat_doc_context, get_chat_doc_context
        await save_chat_doc_context("chat_999", "agromatrix", DOC_A)
        result = await get_chat_doc_context("chat_999", "agromatrix")
    return result


def test_upload_then_question_different_session():
    """
    Сценарій: upload doc (session_1) → питання (session_2) → Stepan бачить doc_id.

    Симулюємо через save_chat_doc_context / get_chat_doc_context:
    - chat_id незмінний, але session_id різний (як у Telegram)
    """
    result = asyncio.run(_scenario_upload_then_question())
    assert result is not None, "Stepan має бачити doc_id після upload"
    assert result["doc_id"] == "parsed_doc_abc123"
    assert "Звіт кукурудза" in result["file_name"]
    assert "прибуток" in result.get("extracted_summary", "").lower()


async def _scenario_agent_isolation():
    stored = {}

    async def mock_upsert(user_id, fact_key, fact_value_json, team_id=None, **kwargs):
        stored[fact_key] = fact_value_json
        return True

    async def mock_get(user_id, fact_key, **kwargs):
        val = stored.get(fact_key)
        return {"fact_value_json": val} if val else None

    mock_client = AsyncMock()
    mock_client.upsert_fact = AsyncMock(side_effect=mock_upsert)
    mock_client.get_fact = AsyncMock(side_effect=mock_get)

    with patch("services.doc_service.memory_client", mock_client):
        from services.doc_service import save_chat_doc_context, get_chat_doc_context
        await save_chat_doc_context("shared_chat", "agromatrix", DOC_A)
        result_helion = await get_chat_doc_context("shared_chat", "helion")
        result_agro = await get_chat_doc_context("shared_chat", "agromatrix")
    return result_helion, result_agro


def test_chat_scoped_key_isolated_between_agents():
    """doc_context_chat:{agent_id}:{chat_id} — різні агенти мають різні ключі."""
    result_helion, result_agro = asyncio.run(_scenario_agent_isolation())
    assert result_helion is None, "Helion не має бачити doc agromatrix"
    assert result_agro is not None, "Agromatrix має бачити свій doc"
    assert result_agro["doc_id"] == "parsed_doc_abc123"


# ── Fix A: dedup tests ───────────────────────────────────────────────────────

async def _scenario_dedup_same_file():
    """Двічі надсилаємо той самий file_unique_id → upsert_fact викликається один раз."""
    stored = {}
    upsert_calls = []

    async def mock_upsert(user_id, fact_key, fact_value_json, team_id=None, **kwargs):
        upsert_calls.append(fact_key)
        stored[fact_key] = fact_value_json
        return True

    async def mock_get(user_id, fact_key, **kwargs):
        val = stored.get(fact_key)
        return {"fact_value_json": val} if val else None

    mock_client = AsyncMock()
    mock_client.upsert_fact = AsyncMock(side_effect=mock_upsert)
    mock_client.get_fact = AsyncMock(side_effect=mock_get)

    with patch("services.doc_service.memory_client", mock_client):
        from services.doc_service import save_chat_doc_context
        # Перший раз — зберегти
        await save_chat_doc_context("chat_dup", "agromatrix", DOC_A)
        # Другий раз той самий file_unique_id → no-op
        await save_chat_doc_context("chat_dup", "agromatrix", DOC_A)

    return upsert_calls


def test_dedup_same_file_unique_id_no_extra_write():
    """Якщо file_unique_id не змінився — upsert_fact викликається лише один раз."""
    calls = asyncio.run(_scenario_dedup_same_file())
    assert len(calls) == 1, f"Expected 1 upsert call, got {len(calls)}: {calls}"


async def _scenario_dedup_different_file():
    """Різний file_unique_id → upsert_fact викликається двічі."""
    stored = {}
    upsert_calls = []

    async def mock_upsert(user_id, fact_key, fact_value_json, team_id=None, **kwargs):
        upsert_calls.append(fact_key)
        stored[fact_key] = fact_value_json
        return True

    async def mock_get(user_id, fact_key, **kwargs):
        val = stored.get(fact_key)
        return {"fact_value_json": val} if val else None

    mock_client = AsyncMock()
    mock_client.upsert_fact = AsyncMock(side_effect=mock_upsert)
    mock_client.get_fact = AsyncMock(side_effect=mock_get)

    doc_b = dict(DOC_A)
    doc_b["file_unique_id"] = "tg_unique_NEW"
    doc_b["doc_id"] = "parsed_doc_def456"

    with patch("services.doc_service.memory_client", mock_client):
        from services.doc_service import save_chat_doc_context
        await save_chat_doc_context("chat_dup2", "agromatrix", DOC_A)
        await save_chat_doc_context("chat_dup2", "agromatrix", doc_b)

    return upsert_calls


def test_dedup_different_file_unique_id_writes_both():
    """Якщо file_unique_id змінився — upsert_fact викликається двічі."""
    calls = asyncio.run(_scenario_dedup_different_file())
    assert len(calls) == 2, f"Expected 2 upsert calls, got {len(calls)}: {calls}"


# ── Fix B: sanitize tests ────────────────────────────────────────────────────

def test_sanitize_summary_removes_rag_prefix():
    """_sanitize_summary видаляє [RAG відповідь по документу]: префікс."""
    from services.doc_service import _sanitize_summary
    raw = "[RAG відповідь по документу]: Прибуток 5 972 016 грн."
    result = _sanitize_summary(raw)
    assert "[RAG" not in result
    assert "Прибуток 5 972 016 грн." in result


def test_sanitize_summary_removes_rag_prefix_variants():
    """_sanitize_summary видаляє різні варіанти [RAG...]: префіксів."""
    from services.doc_service import _sanitize_summary
    variants = [
        "[RAG answer]: some text",
        "[RAG Query Result]: some text",
        "[rag відповідь]: some text",
    ]
    for v in variants:
        result = _sanitize_summary(v)
        assert not result.startswith("[RAG"), f"Not sanitized: {result!r}"
        assert "some text" in result


def test_sanitize_summary_removes_trace_id():
    """_sanitize_summary видаляє trace_id=... артефакти."""
    from services.doc_service import _sanitize_summary
    raw = "Прибуток 5 млн. trace_id=97182015-0c71-4b93-a829-ef5421dea914 Добрива 1.5 млн."
    result = _sanitize_summary(raw)
    assert "trace_id=" not in result
    assert "Прибуток 5 млн." in result
    assert "Добрива 1.5 млн." in result


def test_sanitize_summary_empty_string():
    """_sanitize_summary не ламається на порожньому рядку."""
    from services.doc_service import _sanitize_summary
    assert _sanitize_summary("") == ""
    assert _sanitize_summary(None) is None


def test_save_chat_doc_context_sanitizes_summary():
    """save_chat_doc_context зберігає вже sanitized extracted_summary."""
    mock_client = AsyncMock()
    mock_client.upsert_fact = AsyncMock(return_value=True)
    mock_client.get_fact = AsyncMock(return_value=None)  # no existing → no dedup skip

    dirty_doc = dict(DOC_A)
    dirty_doc["extracted_summary"] = "[RAG відповідь]: Прибуток 5 млн. trace_id=abc-123 Добрива 1 млн."

    with patch("services.doc_service.memory_client", mock_client):
        from services.doc_service import save_chat_doc_context
        run_async(save_chat_doc_context("chat_sanit", "agromatrix", dirty_doc))

    saved = mock_client.upsert_fact.call_args.kwargs["fact_value_json"]
    summary = saved.get("extracted_summary", "")
    assert "[RAG" not in summary, f"RAG prefix leaked: {summary!r}"
    assert "trace_id=" not in summary, f"trace_id leaked: {summary!r}"
    assert "Прибуток 5 млн." in summary