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.")