""" 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