Newer
Older
cortex-hub / ai-hub / tests / api / test_routes.py
import pytest
from unittest.mock import MagicMock, AsyncMock
from fastapi import FastAPI
from fastapi.testclient import TestClient
from sqlalchemy.orm import Session
from datetime import datetime

# Import the dependencies and router factory
from app.core.services import RAGService
from app.api.dependencies import get_db
from app.api.routes import create_api_router
from app.db import models # Import your SQLAlchemy models

@pytest.fixture
def client():
    """Pytest fixture to create a TestClient with a fully mocked environment."""
    test_app = FastAPI()
    mock_rag_service = MagicMock(spec=RAGService)
    mock_db_session = MagicMock(spec=Session)

    def override_get_db():
        yield mock_db_session

    api_router = create_api_router(rag_service=mock_rag_service)
    test_app.dependency_overrides[get_db] = override_get_db
    test_app.include_router(api_router)

    yield TestClient(test_app), mock_rag_service

# --- General Endpoint ---

def test_read_root(client):
    """Tests the root endpoint."""
    test_client, _ = client
    response = test_client.get("/")
    assert response.status_code == 200
    assert response.json() == {"status": "AI Model Hub is running!"}

# --- Session and Chat Endpoints ---

def test_create_session_success(client):
    """Tests successfully creating a new chat session."""
    test_client, mock_rag_service = client
    # Arrange: Mock the service to return a new session object
    mock_session = models.Session(id=1, user_id="test_user", model_name="gemini", title="New Chat", created_at=datetime.now())
    mock_rag_service.create_session.return_value = mock_session
    
    # Act
    response = test_client.post("/sessions", json={"user_id": "test_user", "model": "gemini"})
    
    # Assert
    assert response.status_code == 200
    response_data = response.json()
    assert response_data["id"] == 1
    assert response_data["user_id"] == "test_user"
    assert response_data["model_name"] == "gemini"
    mock_rag_service.create_session.assert_called_once()

def test_chat_in_session_success(client):
    """Tests sending a message in an existing session."""
    test_client, mock_rag_service = client
    # Arrange: Mock the chat service to return a tuple (answer, model_name)
    mock_rag_service.chat_with_rag.return_value = ("Mocked response", "deepseek")
    
    # Act: Send a chat message to a hypothetical session 42
    response = test_client.post("/sessions/42/chat", json={"prompt": "Hello there"})
    
    # Assert
    assert response.status_code == 200
    assert response.json() == {"answer": "Mocked response", "model_used": "deepseek"}
    mock_rag_service.chat_with_rag.assert_called_once()

# --- Document Endpoints ---

def test_add_document_success(client):
    """Tests successfully adding a document."""
    test_client, mock_rag_service = client
    mock_rag_service.add_document.return_value = 123
    doc_payload = {"title": "Test Doc", "text": "Content here"}

    response = test_client.post("/documents", json=doc_payload)

    assert response.status_code == 200
    assert response.json()["message"] == "Document 'Test Doc' added successfully with ID 123"
    mock_rag_service.add_document.assert_called_once()

def test_get_documents_success(client):
    """Tests successfully retrieving a list of all documents."""
    test_client, mock_rag_service = client
    # Arrange: Your mock service should return objects that match the schema attributes
    mock_docs = [
        models.Document(id=1, title="Doc One", status="ready", created_at=datetime.now()),
        models.Document(id=2, title="Doc Two", status="processing", created_at=datetime.now())
    ]
    mock_rag_service.get_all_documents.return_value = mock_docs

    # Act
    response = test_client.get("/documents")
    
    # Assert
    assert response.status_code == 200
    response_data = response.json()
    assert len(response_data["documents"]) == 2
    assert response_data["documents"][0]["title"] == "Doc One"
    mock_rag_service.get_all_documents.assert_called_once()

def test_delete_document_success(client):
    """Tests successfully deleting a document."""
    test_client, mock_rag_service = client
    mock_rag_service.delete_document.return_value = 42

    response = test_client.delete("/documents/42")
    
    assert response.status_code == 200
    assert response.json()["message"] == "Document deleted successfully"
    assert response.json()["document_id"] == 42
    mock_rag_service.delete_document.assert_called_once_with(db=mock_rag_service.delete_document.call_args.kwargs['db'], document_id=42)

def test_delete_document_not_found(client):
    """Tests attempting to delete a document that does not exist."""
    test_client, mock_rag_service = client
    mock_rag_service.delete_document.return_value = None

    response = test_client.delete("/documents/999")
    
    assert response.status_code == 404
    assert response.json()["detail"] == "Document with ID 999 not found."