import os import sys # Add app to path app_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "app")) if app_dir not in sys.path: sys.path.insert(0, app_dir) from fastapi.testclient import TestClient from app import app from core.storage import storage from core.persistent_db import persistent_db client = TestClient(app) def test_full_memory_and_multidoc_isolation(): print("\n--- 1. Testing Session Creation & Isolation ---") res_a = client.post("/chat/new") assert res_a.status_code == 200 chat_a = res_a.json()["chat_id"] res_b = client.post("/chat/new") assert res_b.status_code == 200 chat_b = res_b.json()["chat_b" if "chat_b" in res_b.json() else "chat_id"] print(f"Chat A ID: {chat_a}") print(f"Chat B ID: {chat_b}") assert chat_a != chat_b # Create dummy mock docs in storage doc1_id = "mock-doc-1" doc2_id = "mock-doc-2" mock_tree_1 = [ {"title": "Section 1: Architecture", "page": 1, "text": "Atlas uses tree-based retrieval."}, {"title": "Section 2: Memory", "page": 2, "text": "Session memory isolates short and long context."}, ] mock_tree_2 = [ {"title": "Overview: Financials", "page": 1, "text": "Revenue grew by 24% in Q3."}, ] storage.set_document_ready(doc1_id, tree=mock_tree_1, filename="Architecture_Guide.pdf", page_count=2, section_count=2) storage.set_document_ready(doc2_id, tree=mock_tree_2, filename="Financial_Report_Q3.pdf", page_count=1, section_count=1) print("\n--- 2. Testing Multi-Document Attachment to Chat A ---") attach_res = client.post(f"/chat/{chat_a}/attach-multiple", json={"doc_ids": [doc1_id, doc2_id]}) assert attach_res.status_code == 200 assert len(attach_res.json()["doc_ids"]) == 2 docs_a_res = client.get(f"/chat/{chat_a}/documents") assert docs_a_res.status_code == 200 docs_a = docs_a_res.json()["documents"] print(f"Chat A attached documents count: {len(docs_a)}") assert len(docs_a) == 2 assert {d["doc_id"] for d in docs_a} == {doc1_id, doc2_id} # Verify Chat B has zero attached documents (Strict Session Isolation) docs_b_res = client.get(f"/chat/{chat_b}/documents") assert docs_b_res.status_code == 200 docs_b = docs_b_res.json()["documents"] print(f"Chat B attached documents count: {len(docs_b)}") assert len(docs_b) == 0, "Violation of session isolation: Chat B should have 0 docs" print("\n--- 3. Testing Message Isolation Between Sessions ---") storage.add_message(chat_a, "user", "What is Atlas?") storage.add_message(chat_a, "assistant", "Atlas is a document intelligence platform.") storage.add_message(chat_b, "user", "Hello from user B") storage.add_message(chat_b, "assistant", "Greetings, this is session B.") hist_a = client.get(f"/chat/{chat_a}/history").json()["history"] hist_b = client.get(f"/chat/{chat_b}/history").json()["history"] assert len(hist_a) == 2 assert hist_a[0]["content"] == "What is Atlas?" assert len(hist_b) == 2 assert hist_b[0]["content"] == "Hello from user B" print("Session history isolation verified!") print("\n--- 4. Testing Session Memory Diagnostic Endpoint ---") mem_res = client.get(f"/memory/chat/{chat_a}") assert mem_res.status_code == 200 mem_data = mem_res.json() assert mem_data["chat_id"] == chat_a assert mem_data["recent_turns_count"] == 2 print(f"Memory diagnostic for Chat A: {mem_data['recent_turns_count']} turns recorded.") print("\n--- 5. Testing Detach Single Document ---") detach_res = client.delete(f"/chat/{chat_a}/document/{doc1_id}") assert detach_res.status_code == 200 remaining_docs = client.get(f"/chat/{chat_a}/documents").json()["documents"] assert len(remaining_docs) == 1 assert remaining_docs[0]["doc_id"] == doc2_id print(f"Remaining document in Chat A: {remaining_docs[0]['filename']}") print("\n--- 6. Testing Global Knowledge & Dossier Storage ---") # Save a test dossier persistent_db.save_knowledge_dossier( doc_id=doc1_id, filename="Architecture_Guide.pdf", executive_summary="Explains Atlas tree-based vectorless RAG architecture.", key_entities=["Atlas", "PageIndex", "Groq", "FastAPI"], key_findings=["Tree search outperforms embeddings on dense tabular texts."], topics=["Architecture", "Document Intelligence", "FastAPI"] ) persistent_db.save_global_knowledge_summary( collective_summary="Atlas workspace contains technical documentation on tree retrieval and quarterly financial reports.", total_documents=2, key_insights=["Tree retrieval avoids embedding drift", "Multi-doc synthesis enables cross-sectional queries"] ) g_res = client.get("/memory/knowledge") assert g_res.status_code == 200 g_data = g_res.json() print(f"Global Knowledge Summary: {g_data['collective_summary']}") print(f"Dossiers count: {len(g_data['dossiers'])}") assert len(g_data["dossiers"]) >= 1 assert "Atlas" in g_data["collective_summary"] print("\n ALL BACKEND MULTI-DOC & MEMORY TESTS PASSED PERFECTLY!") if __name__ == "__main__": test_full_memory_and_multidoc_isolation()