Spaces:
Running
Running
File size: 7,485 Bytes
9ad6c13 841b61d 9ad6c13 841b61d 9ad6c13 841b61d 9ad6c13 841b61d 9ad6c13 841b61d 9ad6c13 841b61d 9ad6c13 841b61d 9ad6c13 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 | import os
import sys
from fastapi.testclient import TestClient
# Ensure app path is discoverable
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), "app"))
from app import app
from core.storage import storage
client = TestClient(app)
def test_studio_mindmap_and_quiz_endpoints():
print("\n--- RUNNING PHASE 4 STUDIO BACKEND TEST SUITE ---")
# 1. Create a test document in storage
doc_id = "test-studio-doc-123"
client_id = "00000000-0000-4000-8000-000000000001"
headers = {"X-Client-ID": client_id}
tree = [
{
"node_id": "ch_1",
"title": "1. Introduction & Vectorless Paradigm",
"page_index": 1,
"summary": "Introduction to PageIndex vectorless RAG architecture.",
"nodes": [
{
"node_id": "sec_1_1",
"title": "1.1 The Pitfalls of Vector Chunking",
"page_index": 2,
"summary": "Chunking causes boundary loss and high maintenance costs.",
}
]
},
{
"node_id": "ch_2",
"title": "2. Hierarchical Tree Traversal",
"page_index": 5,
"summary": "Describes DocumentService and LLM reasoning over trees.",
"nodes": [
{
"node_id": "sec_2_1",
"title": "2.1 BM25 Lexical Boosting",
"page_index": 7,
"summary": "Pure Python BM25 scoring in under 1ms.",
}
]
}
]
storage.create_document(
doc_id=doc_id,
filename="Atlas_Architecture.pdf",
status="ready",
owner_id=client_id,
)
storage.set_document_ready(
doc_id=doc_id,
tree=tree,
filename="Atlas_Architecture.pdf",
page_count=10,
section_count=4,
)
# Save a test knowledge dossier
storage.save_knowledge_dossier(
doc_id=doc_id,
filename="Atlas_Architecture.pdf",
executive_summary="Atlas is an enterprise vectorless RAG platform with hierarchical tree search.",
key_entities=["Atlas", "PageIndex", "Groq", "FastAPI"],
key_findings=["Tree search outperforms embeddings on dense technical documents.", "Sub-node pruning preserves exact page citations."],
topics=["Vectorless Architecture", "Hierarchical Traversal", "Lexical BM25"],
owner_id=client_id,
)
# 2. Test GET /studio/{doc_id}/mindmap
print("\n[TEST 1] Testing GET /studio/{doc_id}/mindmap...")
res = client.get(f"/studio/{doc_id}/mindmap", headers=headers)
assert res.status_code == 200, f"Expected 200, got {res.status_code}: {res.text}"
data = res.json()
assert data["doc_id"] == doc_id
assert data["filename"] == "Atlas_Architecture.pdf"
assert "nodes" in data and "edges" in data
nodes = data["nodes"]
edges = data["edges"]
print(f"Mind map generated with {len(nodes)} nodes and {len(edges)} edges.")
# Validate root node
root_node = next((n for n in nodes if n["type"] == "root"), None)
assert root_node is not None, "Missing root node"
assert root_node["section_title"] == "Atlas_Architecture.pdf"
assert root_node["page_number"] == 1
# Validate dossier hub & topic nodes
dossier_hub = next((n for n in nodes if n["type"] == "dossier"), None)
assert dossier_hub is not None, "Missing dossier hub node"
topic_nodes = [n for n in nodes if n["type"] == "topic"]
assert len(topic_nodes) >= 1, "Expected topic nodes from dossier"
# Validate chapters & sections
chapters = [n for n in nodes if n["type"] == "chapter"]
sections = [n for n in nodes if n["type"] == "section"]
assert len(chapters) == 2, f"Expected 2 chapters, got {len(chapters)}"
assert len(sections) == 2, f"Expected 2 sub-sections, got {len(sections)}"
# Validate that every single node has required study studio fields
for n in nodes:
assert "id" in n and n["id"]
assert "label" in n and n["label"]
assert "section_title" in n and n["section_title"]
assert "page_number" in n and n["page_number"] is not None
assert "type" in n
assert "level" in n
# Validate edges
assert len(edges) >= len(nodes) - 1, "Graph should be connected"
print("✓ Test 1 Passed: GET /studio/{doc_id}/mindmap correctly transforms tree & dossier into graph!")
# 3. Test POST /studio/{doc_id}/quiz
print("\n[TEST 2] Testing POST /studio/{doc_id}/quiz...")
quiz_res = client.post(f"/studio/{doc_id}/quiz", headers=headers, json={"num_mcqs": 5, "num_flashcards": 5})
assert quiz_res.status_code == 200, f"Expected 200, got {quiz_res.status_code}: {quiz_res.text}"
quiz_data = quiz_res.json()
assert quiz_data["doc_id"] == doc_id
assert "mcqs" in quiz_data and "flashcards" in quiz_data
mcqs = quiz_data["mcqs"]
flashcards = quiz_data["flashcards"]
print(f"Quiz generated with {len(mcqs)} MCQs and {len(flashcards)} Flashcards.")
assert len(mcqs) == 5, f"Expected exactly 5 MCQs, got {len(mcqs)}"
assert len(flashcards) == 5, f"Expected exactly 5 Flashcards, got {len(flashcards)}"
# Validate MCQs structure & citations
for idx, q in enumerate(mcqs):
assert "question" in q and q["question"], f"MCQ {idx} missing question"
assert "options" in q and len(q["options"]) == 4, f"MCQ {idx} must have 4 choices, got {len(q.get('options', []))}"
assert "answer" in q and q["answer"], f"MCQ {idx} missing answer"
assert "explanation" in q and q["explanation"], f"MCQ {idx} missing explanation"
assert "citation" in q and q["citation"], f"MCQ {idx} missing citation"
assert "(Section:" in q["citation"] and "Page:" in q["citation"], f"MCQ {idx} citation malformed: {q['citation']}"
# Validate Flashcards structure & citations
for idx, c in enumerate(flashcards):
assert "question" in c and c["question"], f"Flashcard {idx} missing question"
assert "answer" in c and c["answer"], f"Flashcard {idx} missing answer"
assert "explanation" in c and c["explanation"], f"Flashcard {idx} missing explanation"
assert "citation" in c and c["citation"], f"Flashcard {idx} missing citation"
assert "(Section:" in c["citation"] and "Page:" in c["citation"], f"Flashcard {idx} citation malformed: {c['citation']}"
print(f"Sample MCQ 1 Citation: {mcqs[0]['citation']}")
print(f"Sample Flashcard 1 Citation: {flashcards[0]['citation']}")
print("✓ Test 2 Passed: POST /studio/{doc_id}/quiz generated 5 MCQs & 5 Flashcards with exact citations!")
# 4. Test 404 for unknown document
print("\n[TEST 3] Testing 404 error handling for non-existent doc...")
err_res = client.get("/studio/non-existent-doc-999/mindmap", headers=headers)
assert err_res.status_code == 404
err_res2 = client.post("/studio/non-existent-doc-999/quiz", headers=headers)
assert err_res2.status_code == 404
print("✓ Test 3 Passed: 404 error handling verified!")
print("\n=======================================================")
print(" ALL 3 PHASE 4 STUDIO BACKEND TESTS PASSED PERFECTLY!")
print("=======================================================\n")
if __name__ == "__main__":
test_studio_mindmap_and_quiz_endpoints()
|