Spaces:
Sleeping
Sleeping
| import pytest | |
| from fastapi.testclient import TestClient | |
| from app.main import app | |
| from app.models.agent import QueryRequest | |
| client = TestClient(app) | |
| class TestAgentEndpoint: | |
| """Integration tests for the agent endpoint.""" | |
| def test_agent_endpoint_selected_text_mode(self): | |
| """Test the agent endpoint with selected text mode.""" | |
| request_data = { | |
| "question": "What does this text mean?", | |
| "selected_text": "The main themes include good vs evil and friendship." | |
| } | |
| response = client.post("/agent/query", json=request_data) | |
| assert response.status_code == 200 | |
| data = response.json() | |
| assert "answer" in data | |
| assert "sources" in data | |
| assert "mode" in data | |
| assert data["mode"] == "selected_text" | |
| def test_agent_endpoint_book_mode(self): | |
| """Test the agent endpoint with book mode (default).""" | |
| request_data = { | |
| "question": "What are the main themes in this book?" | |
| } | |
| response = client.post("/agent/query", json=request_data) | |
| assert response.status_code == 200 | |
| data = response.json() | |
| assert "answer" in data | |
| assert "sources" in data | |
| assert "mode" in data | |
| assert data["mode"] == "book" | |
| def test_agent_endpoint_general_mode(self): | |
| """Test the agent endpoint with general mode.""" | |
| request_data = { | |
| "question": "What is the capital of France?", | |
| "mode": "general" | |
| } | |
| response = client.post("/agent/query", json=request_data) | |
| assert response.status_code == 200 | |
| data = response.json() | |
| assert "answer" in data | |
| assert "sources" in data | |
| assert "mode" in data | |
| assert data["mode"] == "general" | |
| def test_agent_endpoint_invalid_request(self): | |
| """Test the agent endpoint with invalid request data.""" | |
| request_data = { | |
| "question": "" # Empty question should fail validation | |
| } | |
| response = client.post("/agent/query", json=request_data) | |
| # Should return 422 for validation error | |
| assert response.status_code == 422 | |
| def test_agent_endpoint_missing_question(self): | |
| """Test the agent endpoint with missing question.""" | |
| request_data = { | |
| # Missing required question field | |
| } | |
| response = client.post("/agent/query", json=request_data) | |
| # Should return 422 for validation error | |
| assert response.status_code == 422 |