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

# The base URL for the local server started by the run_tests.sh script
BASE_URL = "http://127.0.0.1:8000"

# A common prompt to be used for the tests
TEST_PROMPT = "Explain the theory of relativity in one sentence."

async def test_root_endpoint():
    """
    Tests if the root endpoint is alive and returns the correct status message.
    """
    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.")

async def test_chat_endpoint_deepseek():
    """
    Tests the /chat endpoint using the default 'deepseek' model.
    Verifies a successful response, correct structure, and valid content.
    """
    print("\n--- Running test_chat_endpoint_deepseek ---")
    url = f"{BASE_URL}/chat?model=deepseek"
    payload = {"prompt": TEST_PROMPT}

    async with httpx.AsyncClient(timeout=60.0) as client: # Increased timeout for robustness
        response = await client.post(url, json=payload)

    # 1. Check for a successful response
    assert response.status_code == 200, f"Expected status 200, but got {response.status_code}. Response: {response.text}"
    
    # 2. Check the response structure
    data = response.json()
    assert "response" in data
    assert "model_used" in data
    
    # 3. Validate the content
    assert data["model_used"] == "deepseek"
    assert isinstance(data["response"], str)
    assert len(data["response"]) > 0
    print(f"✅ DeepSeek chat test passed. Response snippet: {data['response'][:80]}...")


async def test_chat_endpoint_gemini():
    """
    Tests the /chat endpoint explicitly requesting the 'gemini' model.
    Verifies a successful response, correct structure, and valid content.
    """
    print("\n--- Running test_chat_endpoint_gemini ---")
    url = f"{BASE_URL}/chat?model=gemini"
    payload = {"prompt": TEST_PROMPT}

    async with httpx.AsyncClient(timeout=60.0) as client: # Increased timeout for robustness
        response = await client.post(url, json=payload)

    # 1. Check for a successful response.
    assert response.status_code == 200, f"Expected status 200, but got {response.status_code}. Response: {response.text}"

    # 2. Check the response structure
    data = response.json()
    assert "response" in data
    assert "model_used" in data

    # 3. Validate the content
    assert data["model_used"] == "gemini"
    assert isinstance(data["response"], str)
    assert len(data["response"]) > 0
    print(f"✅ Gemini chat test passed. Response snippet: {data['response'][:80]}...")


async def test_chat_with_empty_prompt():
    """
    Tests the /chat endpoint's error handling with an empty prompt.
    The Pydantic model should now reject this input with a 422 error
    due to the `min_length` constraint in the ChatRequest model.
    """
    print("\n--- Running test_chat_with_empty_prompt ---")
    url = f"{BASE_URL}/chat"
    payload = {"prompt": ""}

    async with httpx.AsyncClient(timeout=30.0) as client:
        response = await client.post(url, json=payload)

    assert response.status_code == 422
    # This assertion has been updated to match the new Pydantic error message.
    assert "string_too_short" in response.json()["detail"][0]["type"]
    print("✅ Empty prompt test passed.")


async def test_unsupported_model():
    """
    Tests the API's error handling for an invalid model name.
    Expects a 422 Unprocessable Entity error due to Pydantic validation.
    """
    print("\n--- Running test_unsupported_model ---")
    url = f"{BASE_URL}/chat?model=unsupported_model_123"
    payload = {"prompt": TEST_PROMPT}

    async with httpx.AsyncClient() as client:
        response = await client.post(url, json=payload)
    
    assert response.status_code == 422
    assert "Input should be 'deepseek' or 'gemini'" in response.json()["detail"][0]["msg"]
    print("✅ Unsupported model test passed.")


async def test_add_document_success():
    """
    Tests the /document endpoint for successful document ingestion.
    Verifies the response status and the message containing the new document ID.
    """
    print("\n--- Running test_add_document_success ---")
    url = f"{BASE_URL}/document"
    doc_data = {
        "title": "Test Integration Document",
        "text": "This document is for testing the integration endpoint.",
        "source_url": "http://example.com/integration_test"
    }

    async with httpx.AsyncClient(timeout=30.0) as client:
        response = await client.post(url, json=doc_data)

    assert response.status_code == 200
    assert "Document 'Test Integration Document' added successfully" in response.json()["message"]
    print("✅ Add document success test passed.")


async def test_add_document_invalid_data():
    """
    Tests the /document endpoint's error handling for missing required fields.
    Pydantic should return a 422 error for a missing 'title'.
    """
    print("\n--- Running test_add_document_invalid_data ---")
    url = f"{BASE_URL}/document"
    doc_data = {
        "text": "This document is missing a title.",
        "source_url": "http://example.com/invalid_data"
    }

    async with httpx.AsyncClient() as client:
        response = await client.post(url, json=doc_data)

    assert response.status_code == 422
    assert "field required" in response.json()["detail"][0]["msg"].lower()
    print("✅ Add document with invalid data test passed.")