melbinjp/DocQA

1import pytest2from fastapi.testclient import TestClient3from unittest.mock import patch4import os5import datetime6 7# Set a dummy API key for tests8os.environ['GOOGLE_API_KEY'] = 'test-key'9 10# Mock asyncio.create_task BEFORE the app is imported11with patch('asyncio.create_task'):12    from app import app, sessions, _clean_sessions_once13    from user_session import UserSession14    from rag_session import RAGSession15 16# Use a client that handles the lifespan context17@pytest.fixture18def client():19    sessions.clear()20    # Using the 'with' statement ensures that startup and shutdown events are run21    with TestClient(app) as test_client:22        yield test_client23    sessions.clear()24 25def test_create_session(client):26    """Tests that a new user session can be created."""27    response = client.post("/sessions")28    assert response.status_code == 20029    data = response.json()30    assert "session_id" in data31    assert data["session_id"] in sessions32    assert isinstance(sessions[data["session_id"]], UserSession)33 34def test_ingest_into_session(client, mocker):35    """Tests ingesting a document into a created session."""36    session_id = client.post("/sessions").json()["session_id"]37 38    mocker.patch("app.load_source", return_value="Test content")39    response = client.post(40        f"/sessions/{session_id}/ingest",41        files={"file": ("test.txt", b"...", "text/plain")}42    )43    assert response.status_code == 20044    data = response.json()45    assert "doc_id" in data46    assert "num_chunks" in data47    assert data["num_chunks"] == 1  # "Test content" should be a single chunk48 49    user_session = sessions[session_id]50    assert len(user_session.docs) == 151    doc_id = data["doc_id"]52    assert user_session.get_doc(doc_id) is not None53 54def test_query_session(client, mocker):55    """Tests querying documents within a session."""56    session_id = client.post("/sessions").json()["session_id"]57    mocker.patch("app.load_source", return_value="Content A")58    resp_a = client.post(f"/sessions/{session_id}/ingest", files={"file": ("doc_a.txt", b"A", "text/plain")})59    doc_id_a = resp_a.json()["doc_id"]60 61    mocker.patch("app.load_source", return_value="Content B")62    resp_b = client.post(f"/sessions/{session_id}/ingest", files={"file": ("doc_b.txt", b"B", "text/plain")})63    doc_id_b = resp_b.json()["doc_id"]64 65    user_session = sessions[session_id]66    mocker.patch.object(user_session.get_doc(doc_id_a), 'query', return_value=[{"text": "from A", "score": 0.9}])67    mocker.patch.object(user_session.get_doc(doc_id_b), 'query', return_value=[{"text": "from B", "score": 0.8}])68 69    # Mock the async generator70    async def mock_async_gen(*args, **kwargs):71        yield "Final Answer"72    mocker.patch("app.generate_rag_response", side_effect=mock_async_gen)73 74    # Query all docs in session75    response_all = client.post(f"/sessions/{session_id}/query", json={"q": "test"})76    assert response_all.status_code == 20077    assert response_all.json()["answer"] == "Final Answer"78 79    # Query a specific doc in session80    response_specific = client.post(f"/sessions/{session_id}/query", json={"q": "test", "doc_ids": [doc_id_a]})81    assert response_specific.status_code == 20082    assert response_specific.json()["answer"] == "Final Answer"83 84def test_delete_document_from_session(client, mocker):85    """Tests deleting a document from a session."""86    session_id = client.post("/sessions").json()["session_id"]87    mocker.patch("app.load_source", return_value="Test content")88    resp = client.post(f"/sessions/{session_id}/ingest", files={"file": ("test.txt", b"...", "text/plain")})89    doc_id = resp.json()["doc_id"]90 91    assert len(sessions[session_id].docs) == 192 93    response_delete = client.delete(f"/sessions/{session_id}/documents/{doc_id}")94    assert response_delete.status_code == 20495 96    assert len(sessions[session_id].docs) == 097 98def test_session_cleanup_logic():99    """Tests the single-pass cleanup logic directly."""100    sessions.clear()101 102    # We need the model to instantiate the RAGSession103    # In a real test setup, we might mock this, but for now, we rely on the app state104    # This test must run within a context where the app lifespan has started.105    # For direct calling, we can manually set it if needed, or rely on other tests106    # having populated it. Let's ensure it's there.107    if not hasattr(app.state, "embedding_model"):108        from sentence_transformers import SentenceTransformer109        app.state.embedding_model = SentenceTransformer('paraphrase-multilingual-mpnet-base-v2')110 111    fresh_session = UserSession()112    fresh_session.add_doc("doc1", RAGSession(source="fresh.txt", embedding_model=app.state.embedding_model))113    sessions["fresh_session"] = fresh_session114 115    expired_session = UserSession()116    expired_session.add_doc("doc2", RAGSession(source="expired.txt", embedding_model=app.state.embedding_model))117    expired_session.last_accessed = datetime.datetime.now() - datetime.timedelta(minutes=20)118    sessions["expired_session"] = expired_session119 120    assert "fresh_session" in sessions121    assert "expired_session" in sessions122 123    _clean_sessions_once()124 125    assert "fresh_session" in sessions126    assert "expired_session" not in sessions127 128    sessions.clear()129