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 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 response = test_client.post("/sessions", json={"user_id": "test_user", "model": "gemini"}) assert response.status_code == 200 assert response.json()["id"] == 1 mock_rag_service.create_session.assert_called_once() def test_chat_in_session_success(client): """ Tests sending a message in an existing session without specifying a model. It should default to 'deepseek'. """ test_client, mock_rag_service = client mock_rag_service.chat_with_rag = AsyncMock(return_value=("Mocked response", "deepseek")) response = test_client.post("/sessions/42/chat", json={"prompt": "Hello there"}) assert response.status_code == 200 assert response.json() == {"answer": "Mocked response", "model_used": "deepseek"} # Verify that chat_with_rag was called with the default model 'deepseek' mock_rag_service.chat_with_rag.assert_called_once_with( db=mock_rag_service.chat_with_rag.call_args.kwargs['db'], session_id=42, prompt="Hello there", model="deepseek" ) def test_chat_in_session_with_model_switch(client): """ Tests sending a message in an existing session and explicitly switching the model. """ test_client, mock_rag_service = client mock_rag_service.chat_with_rag = AsyncMock(return_value=("Mocked response from Gemini", "gemini")) response = test_client.post("/sessions/42/chat", json={"prompt": "Hello there, Gemini!", "model": "gemini"}) assert response.status_code == 200 assert response.json() == {"answer": "Mocked response from Gemini", "model_used": "gemini"} # Verify that chat_with_rag was called with the specified model 'gemini' mock_rag_service.chat_with_rag.assert_called_once_with( db=mock_rag_service.chat_with_rag.call_args.kwargs['db'], session_id=42, prompt="Hello there, Gemini!", model="gemini" ) def test_get_session_messages_success(client): """Tests retrieving the message history for a session.""" test_client, mock_rag_service = client # Arrange: Mock the service to return a list of message objects mock_history = [ models.Message(sender="user", content="Hello", created_at=datetime.now()), models.Message(sender="assistant", content="Hi there!", created_at=datetime.now()) ] mock_rag_service.get_message_history.return_value = mock_history # Act response = test_client.get("/sessions/123/messages") # Assert assert response.status_code == 200 response_data = response.json() assert response_data["session_id"] == 123 assert len(response_data["messages"]) == 2 assert response_data["messages"][0]["sender"] == "user" assert response_data["messages"][1]["content"] == "Hi there!" mock_rag_service.get_message_history.assert_called_once_with(db=mock_rag_service.get_message_history.call_args.kwargs['db'], session_id=123) def test_get_session_messages_not_found(client): """Tests retrieving messages for a session that does not exist.""" test_client, mock_rag_service = client # Arrange: Mock the service to return None, indicating the session wasn't found mock_rag_service.get_message_history.return_value = None # Act response = test_client.get("/sessions/999/messages") # Assert assert response.status_code == 404 assert response.json()["detail"] == "Session with ID 999 not found." # --- Document Endpoints --- # (These tests are unchanged) def test_add_document_success(client): 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" def test_get_documents_success(client): test_client, mock_rag_service = client 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 response = test_client.get("/documents") assert response.status_code == 200 assert len(response.json()["documents"]) == 2 def test_delete_document_success(client): 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()["document_id"] == 42 def test_delete_document_not_found(client): 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