Newer
Older
cortex-hub / ai-hub / integration_tests / test_sessions_api.py
import pytest

# Test prompts and data
CONTEXT_PROMPT = "Who is the CEO of Microsoft?"
FOLLOW_UP_PROMPT = "When was he born?"
RAG_DOC_TITLE = "Fictional Company History"
RAG_DOC_TEXT = "The company AlphaCorp was founded in 2021 by Jane Doe. Their primary product is a smart home device called 'Nexus'."
RAG_PROMPT = "Who founded AlphaCorp and what is their main product?"

@pytest.mark.asyncio
async def test_chat_in_session_lifecycle(http_client):
    """
    Tests a full session lifecycle from creation to conversational memory.
    This test is a single, sequential unit.
    """
    print("\n--- Running test_chat_in_session_lifecycle ---")
    
    # 1. Create a new session with a trailing slash
    payload = {"user_id": "integration_tester_lifecycle", "provider_name": "deepseek"}
    response = await http_client.post("/sessions/", json=payload)
    assert response.status_code == 200
    session_id = response.json()["id"]
    print(f"✅ Session created successfully with ID: {session_id}")
    
    # 2. First chat turn to establish context
    chat_payload_1 = {"prompt": CONTEXT_PROMPT}
    response_1 = await http_client.post(f"/sessions/{session_id}/chat", json=chat_payload_1)
    assert response_1.status_code == 200
    assert "Satya Nadella" in response_1.json()["answer"]
    assert response_1.json()["provider_used"] == "deepseek"
    print("✅ Chat Turn 1 (context) test passed.")
    
    # 3. Second chat turn (follow-up) to test conversational memory
    chat_payload_2 = {"prompt": FOLLOW_UP_PROMPT}
    response_2 = await http_client.post(f"/sessions/{session_id}/chat", json=chat_payload_2)
    assert response_2.status_code == 200
    assert "1967" in response_2.json()["answer"]
    assert response_2.json()["provider_used"] == "deepseek"
    print("✅ Chat Turn 2 (follow-up) test passed.")
    
    # 4. Cleanup (optional, but good practice if not using a test database that resets)
    # The session data would typically be cleaned up by the database teardown.

@pytest.mark.asyncio
async def test_chat_with_model_switch(http_client, session_id):
    """Tests switching models within an existing session."""
    print("\n--- Running test_chat_with_model_switch ---")
    
    # Send a message to the new session with a different model
    payload_gemini = {"prompt": "What is the capital of France?", "provider_name": "gemini"}
    response_gemini = await http_client.post(f"/sessions/{session_id}/chat", json=payload_gemini)
    assert response_gemini.status_code == 200
    assert "Paris" in response_gemini.json()["answer"]
    assert response_gemini.json()["provider_used"] == "gemini"
    print("✅ Chat (Model Switch to Gemini) test passed.")

    # Switch back to the original model
    payload_deepseek = {"prompt": "What is the largest ocean?", "provider_name": "deepseek"}
    response_deepseek = await http_client.post(f"/sessions/{session_id}/chat", json=payload_deepseek)
    assert response_deepseek.status_code == 200
    assert "Pacific Ocean" in response_deepseek.json()["answer"]
    assert response_deepseek.json()["provider_used"] == "deepseek"
    print("✅ Chat (Model Switch back to DeepSeek) test passed.")

@pytest.mark.asyncio
async def test_chat_with_document_retrieval(http_client):
    """
    Tests injecting a document and using it for retrieval-augmented generation.
    This test creates its own session and document for isolation.
    """
    print("\n--- Running test_chat_with_document_retrieval ---")
    
    # Create a new session for this RAG test
    session_response = await http_client.post("/sessions/", json={"user_id": "rag_tester", "provider_name": "deepseek"})
    assert session_response.status_code == 200
    rag_session_id = session_response.json()["id"]
    
    # Add a new document with specific content for retrieval
    doc_data = {"title": RAG_DOC_TITLE, "text": RAG_DOC_TEXT}
    add_doc_response = await http_client.post("/documents/", json=doc_data)
    assert add_doc_response.status_code == 200
    try:
        message = add_doc_response.json().get("message", "")
        rag_document_id = int(message.split(" with ID ")[-1])
        print(f"Document for RAG created with ID: {rag_document_id}")
    except (ValueError, IndexError):
        pytest.fail("Could not parse document ID from response message.")
    
    try:
        chat_payload = {
            "prompt": RAG_PROMPT,
            "document_id": rag_document_id,
            "provider_name": "deepseek",
            "load_faiss_retriever": True
        }
        chat_response = await http_client.post(f"/sessions/{rag_session_id}/chat", json=chat_payload)
        
        # --- MODIFICATION START ---
        # If a 500 error occurs, print the detailed response text.
        if chat_response.status_code != 200:
            print(f"❌ Test Failed! Received status code: {chat_response.status_code}")
            print("--- Response Body (for debugging) ---")
            print(chat_response.text)
            print("---------------------------------------")
        
        assert chat_response.status_code == 200
        # --- MODIFICATION END ---
        
        chat_data = chat_response.json()
        assert "Jane Doe" in chat_data["answer"]
        assert "Nexus" in chat_data["answer"]
        print("✅ Chat with document retrieval test passed.")
    finally:
        # Clean up the document after the test
        delete_response = await http_client.delete(f"/documents/{rag_document_id}")
        assert delete_response.status_code == 200
        print(f"Document {rag_document_id} deleted successfully.")


# --- New Session Management Integration Tests ---

@pytest.mark.asyncio
async def test_create_session_with_feature_name(http_client):
    """
    Tests that the feature_name field is accepted on session creation and returned
    in the response. This validates the DB column was added correctly.
    """
    print("\n--- Running test_create_session_with_feature_name ---")
    payload = {
        "user_id": "integration_tester_feature",
        "provider_name": "deepseek",
        "feature_name": "coding_assistant"
    }
    response = await http_client.post("/sessions/", json=payload)
    assert response.status_code == 200, f"Unexpected status: {response.status_code} — {response.text}"
    data = response.json()
    session_id = data["id"]
    assert data["feature_name"] == "coding_assistant"
    print(f"✅ Session created with feature_name='coding_assistant', ID={session_id}")


@pytest.mark.asyncio
async def test_list_sessions_by_feature(http_client):
    """
    Tests that GET /sessions/?user_id=...&feature_name=... returns only sessions
    for the specified feature, not ones from another feature.
    """
    print("\n--- Running test_list_sessions_by_feature ---")
    user_id = "integration_tester_list"

    # Create a coding_assistant session
    r1 = await http_client.post("/sessions/", json={
        "user_id": user_id, "provider_name": "deepseek", "feature_name": "coding_assistant"
    })
    assert r1.status_code == 200
    coding_session_id = r1.json()["id"]

    # Create a voice_chat session for the same user
    r2 = await http_client.post("/sessions/", json={
        "user_id": user_id, "provider_name": "deepseek", "feature_name": "voice_chat"
    })
    assert r2.status_code == 200
    voice_session_id = r2.json()["id"]

    # List only coding_assistant sessions
    list_resp = await http_client.get(f"/sessions/?user_id={user_id}&feature_name=coding_assistant")
    assert list_resp.status_code == 200
    sessions = list_resp.json()

    ids = [s["id"] for s in sessions]
    assert coding_session_id in ids, "coding_assistant session should be in the list"
    assert voice_session_id not in ids, "voice_chat session should NOT appear in coding_assistant list"
    print(f"✅ Session list isolation test passed. coding={coding_session_id}, voice={voice_session_id}")


@pytest.mark.asyncio
async def test_delete_single_session(http_client):
    """
    Tests that DELETE /sessions/{session_id} archives (soft-deletes) the session so it
    no longer appears in the GET /sessions/ list.
    """
    print("\n--- Running test_delete_single_session ---")
    user_id = "integration_tester_delete"

    create_resp = await http_client.post("/sessions/", json={
        "user_id": user_id, "provider_name": "deepseek", "feature_name": "coding_assistant"
    })
    assert create_resp.status_code == 200
    session_id = create_resp.json()["id"]

    # Delete the session
    delete_resp = await http_client.delete(f"/sessions/{session_id}")
    assert delete_resp.status_code == 200
    assert "deleted" in delete_resp.json().get("message", "").lower()

    # Verify the session is no longer returned in the list
    list_resp = await http_client.get(f"/sessions/?user_id={user_id}&feature_name=coding_assistant")
    assert list_resp.status_code == 200
    ids = [s["id"] for s in list_resp.json()]
    assert session_id not in ids, "Deleted session should not appear in the list"
    print(f"✅ Single session delete test passed. Deleted session ID={session_id}")


@pytest.mark.asyncio
async def test_delete_session_not_found(http_client):
    """Tests that deleting a non-existent session returns a 404."""
    print("\n--- Running test_delete_session_not_found ---")
    response = await http_client.delete("/sessions/999999")
    assert response.status_code == 404
    print("✅ Delete non-existent session returns 404 as expected.")


@pytest.mark.asyncio
async def test_delete_all_sessions_for_feature(http_client):
    """
    Tests that DELETE /sessions/?user_id=...&feature_name=... wipes all sessions for
    that feature and they are no longer listed.
    """
    print("\n--- Running test_delete_all_sessions_for_feature ---")
    user_id = "integration_tester_delete_all"

    # Create two sessions for coding_assistant
    for _ in range(2):
        r = await http_client.post("/sessions/", json={
            "user_id": user_id, "provider_name": "deepseek", "feature_name": "coding_assistant"
        })
        assert r.status_code == 200

    # Also create a voice_chat session to ensure it is NOT deleted
    voice_resp = await http_client.post("/sessions/", json={
        "user_id": user_id, "provider_name": "deepseek", "feature_name": "voice_chat"
    })
    assert voice_resp.status_code == 200
    voice_session_id = voice_resp.json()["id"]

    # Bulk delete coding_assistant sessions
    delete_resp = await http_client.delete(
        f"/sessions/?user_id={user_id}&feature_name=coding_assistant"
    )
    assert delete_resp.status_code == 200
    assert "deleted" in delete_resp.json().get("message", "").lower()

    # Confirm coding sessions are gone
    coding_list = await http_client.get(f"/sessions/?user_id={user_id}&feature_name=coding_assistant")
    assert coding_list.status_code == 200
    assert len(coding_list.json()) == 0, "All coding_assistant sessions should be deleted"

    # Confirm voice_chat session is still present
    voice_list = await http_client.get(f"/sessions/?user_id={user_id}&feature_name=voice_chat")
    assert voice_list.status_code == 200
    voice_ids = [s["id"] for s in voice_list.json()]
    assert voice_session_id in voice_ids, "voice_chat session should be unaffected"
    print("✅ Bulk delete by feature test passed. Voice session preserved.")