test_production_async_workflow.py

import os
import sys
import time
import io
import asyncio
import fitz
import httpx

# Ensure app is discoverable
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from app.app import app
from app.core.storage import storage
from app.services.document_service import DocumentService


def create_test_pdf(num_pages: int = 3) -> bytes:
    """Create a minimal valid multi-page PDF in memory."""
    doc = fitz.open()
    for i in range(num_pages):
        page = doc.new_page(width=595, height=842)
        page.insert_text((50, 72), f"Page {i + 1} Heading: Introduction to Section {i + 1}", fontsize=14)
        page.insert_text((50, 120), f"This is page {i + 1} content explaining architecture principles and design.", fontsize=11)
    buf = io.BytesIO()
    doc.save(buf)
    doc.close()
    return buf.getvalue()


async def run_tests():
    results = {}
    transport = httpx.ASGITransport(app=app)

    TEST_CLIENT_ID = "00000000-0000-4000-8000-000000000001"
    async with httpx.AsyncClient(transport=transport, base_url="http://test", headers={"X-Client-ID": TEST_CLIENT_ID}) as client:
        print("\n=======================================================")
        print("STARTING PRODUCTION ASYNC INDEXING TEST SUITE")
        print("=======================================================\n")

        # ─────────────────────────────────────────────────────────────────
        # Scenario A: Fresh PDF Upload
        # ─────────────────────────────────────────────────────────────────
        print("--- [TEST A] Fresh PDF Upload ---")
        pdf_bytes = create_test_pdf(num_pages=4)
        upload_t0 = time.monotonic()
        res_upload = await client.post(
            "/upload",
            files={"file": ("test_fresh.pdf", pdf_bytes, "application/pdf")},
        )
        upload_duration = time.monotonic() - upload_t0
        assert res_upload.status_code == 202, f"Expected 202, got {res_upload.status_code}"
        data_a = res_upload.json()
        assert data_a["status"] == "indexing", f"Expected indexing, got {data_a['status']}"
        assert "doc_id" in data_a and "document_id" in data_a
        assert data_a["doc_id"] == data_a["document_id"]
        doc_id_a = data_a["doc_id"]
        # Verify non-blocking: upload response returned quickly
        assert upload_duration < 3.0, f"Upload took too long: {upload_duration}s"
        print(f"✓ Test A Passed: Returned HTTP 202 in {upload_duration:.3f}s with doc_id={doc_id_a}")
        results["A_fresh_upload"] = True

        # Verify immediate status is indexing
        res_st_init = await client.get(f"/documents/{doc_id_a}/status")
        assert res_st_init.status_code == 200
        st_init = res_st_init.json()
        assert st_init["status"] in ("indexing", "ready")
        print(f"✓ Immediate status check: status={st_init['status']}, stage={st_init['current_stage']}, progress={st_init['progress']}%")

        # ─────────────────────────────────────────────────────────────────
        # Scenario B: Duplicate PDF Upload While Indexing
        # ─────────────────────────────────────────────────────────────────
        print("\n--- [TEST B] Duplicate PDF Upload While Indexing ---")
        res_dup_indexing = await client.post(
            "/upload",
            files={"file": ("test_fresh.pdf", pdf_bytes, "application/pdf")},
        )
        assert res_dup_indexing.status_code in (200, 202)
        data_b = res_dup_indexing.json()
        assert data_b["doc_id"] == doc_id_a, "Duplicate upload should reuse same doc_id"
        print(f"✓ Test B Passed: Reused doc_id={data_b['doc_id']} without duplicate job submission (status={data_b['status']})")
        results["B_duplicate_indexing"] = True

        # ─────────────────────────────────────────────────────────────────
        # Scenario C & F: Polling Status & Browser Refresh Simulation
        # ─────────────────────────────────────────────────────────────────
        print("\n--- [TEST C & F] Status Endpoint & Polling (Browser Refresh Simulation) ---")
        # Poll status endpoint every 0.5s until ready or max 25s
        poll_start = time.monotonic()
        is_ready = False
        last_status = None
        while time.monotonic() - poll_start < 25:
            res_st = await client.get(f"/documents/{doc_id_a}/status")
            assert res_st.status_code == 200
            last_status = res_st.json()
            assert last_status["document_id"] == doc_id_a
            assert "current_stage" in last_status
            assert "progress" in last_status
            print(f"  Polling {time.monotonic() - poll_start:.1f}s: status={last_status['status']} progress={last_status['progress']}% stage={last_status['current_stage']}")
            if last_status["status"] == "ready":
                is_ready = True
                break
            await asyncio.sleep(0.5)

        assert is_ready, f"Document did not become ready in time: {last_status}"
        assert last_status["progress"] == 100
        assert last_status["current_stage"] == "Ready"
        print(f"✓ Test C & F Passed: Polled status to completion. Final status: {last_status['status']}, stage: {last_status['current_stage']}, pages: {last_status['page_count']}, sections: {last_status['section_count']}")
        results["C_slow_response_poll"] = True
        results["F_browser_refresh_polling"] = True

        # ─────────────────────────────────────────────────────────────────
        # Metadata Accuracy Verification (pages vs sections)
        # ─────────────────────────────────────────────────────────────────
        print("\n--- [METADATA CHECK] Accurate page_count vs section_count ---")
        assert last_status["page_count"] == 4, f"Expected 4 pages, got {last_status['page_count']}"
        assert last_status["section_count"] >= 1, f"Expected at least 1 section, got {last_status['section_count']}"
        print(f"✓ Metadata Check Passed: pages={last_status['page_count']} (physical pages), sections={last_status['section_count']}")

        # ─────────────────────────────────────────────────────────────────
        # Duplicate Upload When Ready (Instant Cache Hit)
        # ─────────────────────────────────────────────────────────────────
        print("\n--- [DEDUPLICATION CHECK] Upload identical PDF when ready ---")
        res_ready_dup = await client.post(
            "/upload",
            files={"file": ("test_fresh.pdf", pdf_bytes, "application/pdf")},
        )
        assert res_ready_dup.status_code == 200
        data_ready_dup = res_ready_dup.json()
        assert data_ready_dup["status"] == "ready"
        assert data_ready_dup["doc_id"] == doc_id_a
        print(f"✓ Deduplication Passed: Instant reuse of cached index for doc_id={doc_id_a}")

        # ─────────────────────────────────────────────────────────────────
        # Scenario D: Incomplete / Invalid Tree Validation
        # ─────────────────────────────────────────────────────────────────
        print("\n--- [TEST D] Incomplete / Empty Tree Validation ---")
        assert DocumentService._is_valid_tree(None) is False
        assert DocumentService._is_valid_tree([]) is False
        assert DocumentService._is_valid_tree([{}]) is False
        assert DocumentService._is_valid_tree([{"title": "", "text": "", "summary": ""}]) is False
        assert DocumentService._is_valid_tree([{"title": "Section 1", "text": "Content"}]) is True
        print("✓ Test D Passed: Tree validation correctly rejects empty/incomplete trees and accepts valid ones")
        results["D_incomplete_tree_validation"] = True

        # ─────────────────────────────────────────────────────────────────
        # Scenario E: PyMuPDF Fallback Execution
        # ─────────────────────────────────────────────────────────────────
        print("\n--- [TEST E] PyMuPDF Fallback Execution ---")
        fallback_service = DocumentService(api_key="")
        fallback_service.api_keys = []
        
        temp_pdf_path = "data/test_fallback.pdf"
        os.makedirs("data", exist_ok=True)
        with open(temp_pdf_path, "wb") as f:
            f.write(create_test_pdf(num_pages=5))

        fb_result = fallback_service.process_pdf(temp_pdf_path, document_id="test_fallback_doc")
        assert fb_result["indexing_method"] == "pymupdf"
        assert fb_result["page_count"] == 5
        assert fb_result["section_count"] >= 1
        assert fb_result["duration"] > 0
        assert fb_result["failure_reason"] is not None
        assert isinstance(fb_result["tree"], list) and len(fb_result["tree"]) > 0
        print(f"✓ Test E Passed: PyMuPDF fallback returned method={fb_result['indexing_method']}, pages={fb_result['page_count']}, sections={fb_result['section_count']}, duration={fb_result['duration']}s")
        results["E_pymupdf_fallback"] = True

        # ─────────────────────────────────────────────────────────────────
        # Scenario G: Retry While Indexing
        # ─────────────────────────────────────────────────────────────────
        print("\n--- [TEST G] Retry While Indexing ---")
        indexing_doc_id = storage.create_document(filename="test_active.pdf", status="indexing", owner_id=TEST_CLIENT_ID)
        storage.update_document_progress(indexing_doc_id, status="indexing", current_stage="Building document structure", progress=40)
        
        res_retry_active = await client.post(f"/documents/{indexing_doc_id}/retry")
        assert res_retry_active.status_code == 202
        assert res_retry_active.json()["message"] == "Document is already indexing"
        print("✓ Test G Passed: Retry while indexing safely prevented duplicate execution")
        results["G_retry_while_indexing"] = True

        # ─────────────────────────────────────────────────────────────────
        # Scenario H: Retry After Failed Indexing
        # ─────────────────────────────────────────────────────────────────
        print("\n--- [TEST H] Retry After Failed Indexing ---")
        failed_doc_id = storage.create_document(
            filename="test_failed.pdf",
            status="failed",
            file_path=temp_pdf_path,
            owner_id=TEST_CLIENT_ID,
        )
        storage.set_document_error(failed_doc_id, "Simulated network timeout")
        
        res_failed_st = await client.get(f"/documents/{failed_doc_id}/status")
        assert res_failed_st.json()["status"] == "failed"
        assert res_failed_st.json()["error"] == "Simulated network timeout"

        res_retry = await client.post(f"/documents/{failed_doc_id}/retry")
        assert res_retry.status_code == 202
        assert res_retry.json()["status"] == "indexing"

        # Poll until recovered to ready
        retry_poll_t0 = time.monotonic()
        retry_ready = False
        while time.monotonic() - retry_poll_t0 < 30:
            st = (await client.get(f"/documents/{failed_doc_id}/status")).json()
            if st["status"] == "ready":
                retry_ready = True
                break
            await asyncio.sleep(0.3)

        assert retry_ready, f"Failed document did not recover after retry: {st}"
        print(f"✓ Test H Passed: Successfully retried failed document {failed_doc_id} -> status: {st['status']}")
        results["H_retry_after_failed"] = True

        # ─────────────────────────────────────────────────────────────────
        # Scenario I: Ask Question After Indexing is Ready
        # ─────────────────────────────────────────────────────────────────
        print("\n--- [TEST I] Ask Question After Indexing is Ready ---")
        chat_res = await client.post("/chat/new")
        assert chat_res.status_code == 200
        chat_id = chat_res.json()["chat_id"]

        attach_res = await client.post(f"/chat/{chat_id}/attach", json={"doc_id": doc_id_a})
        assert attach_res.status_code == 200

        ask_res = await client.post(
            "/ask",
            json={"chat_id": chat_id, "query": "What is the content on page 1?"},
        )
        assert ask_res.status_code == 200
        answer_text = ask_res.text
        assert len(answer_text) > 0
        print(f"✓ Test I Passed: Received grounded answer for chat_id={chat_id}, doc_id={doc_id_a}. Answer length: {len(answer_text)} chars")
        results["I_ask_question_ready"] = True

        print("\n=======================================================")
        print("ALL 9 TEST SCENARIOS PASSED SUCCESSFULLY!")
        print("=======================================================\n")
        return results


if __name__ == "__main__":
    asyncio.run(run_tests())