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

# The base URL for the local server
BASE_URL = "http://127.0.0.1:8000"
TEST_PROMPT = "Explain the theory of relativity in one sentence."

# Global variables to pass state between sequential tests
created_document_id = None
created_session_id = None

async def test_root_endpoint():
    """Tests if the root endpoint is alive."""
    print("\n--- Running test_root_endpoint ---")
    async with httpx.AsyncClient() as client:
        response = await client.get(f"{BASE_URL}/")
        assert response.status_code == 200
        assert response.json() == {"status": "AI Model Hub is running!"}
    print("✅ Root endpoint test passed.")

# --- Session and Chat Lifecycle Tests ---

async def test_create_session():
    """
    Tests creating a new chat session and saves the ID for the next test.
    """
    global created_session_id
    print("\n--- Running test_create_session ---")
    url = f"{BASE_URL}/sessions"
    payload = {"user_id": "integration_tester", "model": "deepseek"}
    
    async with httpx.AsyncClient() as client:
        response = await client.post(url, json=payload)

    assert response.status_code == 200, f"Failed to create session. Response: {response.text}"
    response_data = response.json()
    assert "id" in response_data
    assert response_data["user_id"] == "integration_tester"
    assert response_data["model_name"] == "deepseek"
    
    created_session_id = response_data["id"]
    print(f"✅ Session created successfully with ID: {created_session_id}")

async def test_chat_in_session():
    """
    Tests sending a message within the session created by the previous test.
    """
    print("\n--- Running test_chat_in_session ---")
    assert created_session_id is not None, "Session ID was not set by the create_session test."
    
    url = f"{BASE_URL}/sessions/{created_session_id}/chat"
    payload = {"prompt": TEST_PROMPT}
    
    async with httpx.AsyncClient(timeout=60.0) as client:
        response = await client.post(url, json=payload)
        
    assert response.status_code == 200, f"Chat request failed. Response: {response.text}"
    response_data = response.json()
    assert "answer" in response_data
    assert len(response_data["answer"]) > 0
    assert response_data["model_used"] == "deepseek"
    print("✅ Chat in session test passed.")

# --- Document Management Lifecycle Tests ---

async def test_add_document_for_lifecycle():
    """
    Adds a document and saves its ID to be used by the list and delete tests.
    """
    global created_document_id
    print("\n--- Running test_add_document (for lifecycle) ---")
    url = f"{BASE_URL}/documents"
    doc_data = {"title": "Lifecycle Test Doc", "text": "This doc will be listed and deleted."}
    
    async with httpx.AsyncClient(timeout=30.0) as client:
        response = await client.post(url, json=doc_data)

    assert response.status_code == 200, f"Failed to add document. Response: {response.text}"
    response_data = response.json()
    message = response_data.get("message", "")
    assert "added successfully with ID" in message

    try:
        created_document_id = int(message.split(" with ID ")[-1])
    except (ValueError, IndexError):
        pytest.fail("Could not parse document ID from response message.")
    
    print(f"✅ Document for lifecycle test created with ID: {created_document_id}")

async def test_list_documents():
    """
    Tests listing documents to ensure the previously created one appears.
    """
    print("\n--- Running test_list_documents ---")
    assert created_document_id is not None, "Document ID was not set by the add test."
    
    url = f"{BASE_URL}/documents"
    async with httpx.AsyncClient() as client:
        response = await client.get(url)

    assert response.status_code == 200
    response_data = response.json()
    assert "documents" in response_data
    
    ids_in_response = {doc["id"] for doc in response_data["documents"]}
    assert created_document_id in ids_in_response
    print("✅ Document list test passed.")

async def test_delete_document():
    """
    Tests deleting the document created at the start of the lifecycle.
    """
    print("\n--- Running test_delete_document ---")
    assert created_document_id is not None, "Document ID was not set by the add test."

    url = f"{BASE_URL}/documents/{created_document_id}"
    async with httpx.AsyncClient() as client:
        response = await client.delete(url)
    
    assert response.status_code == 200
    response_data = response.json()
    assert response_data["message"] == "Document deleted successfully"
    assert response_data["document_id"] == created_document_id
    print("✅ Document delete test passed.")