| """ |
| tests/integration/test_query_api.py |
| FastAPI integration tests — agent graph is fully mocked so no LLM/DB calls happen. |
| Uses TestClient (synchronous ASGI runner). |
| """ |
|
|
| import json |
| import uuid |
| import pytest |
| from unittest.mock import patch, MagicMock |
|
|
|
|
| def _make_final_state(session_id: str) -> dict: |
| return { |
| "session_id": session_id, |
| "user_id": "test", |
| "user_query": "top products", |
| "connector_id": "neon:public", |
| "intent": "sql", |
| "query_plan": {}, |
| "relevant_tables": [], |
| "schema_context": "", |
| "memory_context": "", |
| "generated_code": "SELECT product, SUM(amount) FROM orders GROUP BY product", |
| "code_type": "sql", |
| "sql_dialect": "postgres", |
| "execution_result": [{"product": "Widget", "total": 299.97}], |
| "execution_error": None, |
| "from_cache": False, |
| "error_class": None, |
| "correction_attempts": 0, |
| "max_corrections": 3, |
| "insight_text": "Widget is the top product with $300 in revenue.", |
| "chart_spec": { |
| "type": "bar", |
| "plotly_json": { |
| "data": [{"type": "bar", "x": ["Widget"], "y": [299.97]}], |
| "layout": {"title": "top products"}, |
| }, |
| }, |
| "history_id": str(uuid.uuid4()), |
| "latency_ms": 850, |
| "stream_tokens": [], |
| } |
|
|
|
|
| @pytest.fixture |
| def mocked_graph(session_id): |
| """Patch get_graph() to return a mock that returns controlled state.""" |
| state = _make_final_state(session_id) |
| mock_graph = MagicMock() |
| mock_graph.invoke.return_value = state |
| with patch("api.routers.query.get_graph", return_value=mock_graph): |
| yield mock_graph, state |
|
|
|
|
| @pytest.mark.integration |
| class TestQueryRunEndpoint: |
| def test_returns_200(self, api_client, mocked_graph, session_id): |
| resp = api_client.post("/api/query/run", json={ |
| "user_query": "top products", |
| "connector_id": "neon:public", |
| "session_id": session_id, |
| }) |
| assert resp.status_code == 200 |
|
|
| def test_response_has_required_fields(self, api_client, mocked_graph, session_id): |
| resp = api_client.post("/api/query/run", json={ |
| "user_query": "top products", |
| "connector_id": "neon:public", |
| "session_id": session_id, |
| }) |
| body = resp.json() |
| required = ["session_id", "intent", "generated_code", "execution_result", |
| "insight_text", "chart_spec", "from_cache", "latency_ms"] |
| for field in required: |
| assert field in body, f"Missing field: {field}" |
|
|
| def test_insight_text_returned(self, api_client, mocked_graph, session_id): |
| resp = api_client.post("/api/query/run", json={ |
| "user_query": "top products", |
| "connector_id": "neon:public", |
| "session_id": session_id, |
| }) |
| assert resp.json()["insight_text"] == "Widget is the top product with $300 in revenue." |
|
|
| def test_chart_spec_returned(self, api_client, mocked_graph, session_id): |
| resp = api_client.post("/api/query/run", json={ |
| "user_query": "top products", |
| "connector_id": "neon:public", |
| "session_id": session_id, |
| }) |
| chart = resp.json()["chart_spec"] |
| assert chart is not None |
| assert chart["type"] == "bar" |
|
|
| def test_empty_query_rejected(self, api_client): |
| resp = api_client.post("/api/query/run", json={ |
| "user_query": "", |
| "connector_id": "neon:public", |
| }) |
| assert resp.status_code == 422 |
|
|
| def test_missing_connector_id_rejected(self, api_client): |
| resp = api_client.post("/api/query/run", json={"user_query": "test"}) |
| assert resp.status_code == 422 |
|
|
| def test_session_id_auto_generated_when_absent(self, api_client, mocked_graph): |
| resp = api_client.post("/api/query/run", json={ |
| "user_query": "top products", |
| "connector_id": "neon:public", |
| }) |
| assert resp.status_code == 200 |
| assert resp.json()["session_id"] |
|
|
| def test_graph_invoke_called_once(self, api_client, mocked_graph, session_id): |
| mock_graph, _ = mocked_graph |
| api_client.post("/api/query/run", json={ |
| "user_query": "top products", |
| "connector_id": "neon:public", |
| "session_id": session_id, |
| }) |
| mock_graph.invoke.assert_called_once() |
|
|
| def test_graph_exception_returns_500(self, api_client, session_id): |
| mock_graph = MagicMock() |
| mock_graph.invoke.side_effect = RuntimeError("LLM timeout") |
| with patch("api.routers.query.get_graph", return_value=mock_graph): |
| resp = api_client.post("/api/query/run", json={ |
| "user_query": "top products", |
| "connector_id": "neon:public", |
| "session_id": session_id, |
| }) |
| assert resp.status_code == 500 |
|
|
|
|
| @pytest.mark.integration |
| class TestHealthEndpoint: |
| def test_health_returns_ok(self, api_client): |
| resp = api_client.get("/health") |
| assert resp.status_code == 200 |
| assert resp.json() == {"status": "ok", "version": "2.0.0"} |
|
|
|
|
| @pytest.mark.integration |
| class TestStreamEndpoint: |
| def test_stream_returns_200(self, api_client, mocked_graph, session_id): |
| resp = api_client.post("/api/query/stream", json={ |
| "user_query": "top products", |
| "connector_id": "neon:public", |
| "session_id": session_id, |
| }) |
| assert resp.status_code == 200 |
|
|
| def test_stream_content_type_is_event_stream(self, api_client, mocked_graph, session_id): |
| resp = api_client.post("/api/query/stream", json={ |
| "user_query": "top products", |
| "connector_id": "neon:public", |
| "session_id": session_id, |
| }) |
| assert "text/event-stream" in resp.headers.get("content-type", "") |
|
|
| def test_stream_contains_done_event(self, api_client, mocked_graph, session_id): |
| resp = api_client.post("/api/query/stream", json={ |
| "user_query": "top products", |
| "connector_id": "neon:public", |
| "session_id": session_id, |
| }) |
| raw = resp.text |
| events = [ |
| json.loads(line[len("data: "):]) |
| for line in raw.split("\n") |
| if line.startswith("data: ") |
| ] |
| done_events = [e for e in events if e.get("done")] |
| assert len(done_events) == 1 |
|
|
| def test_stream_emits_token_events(self, api_client, mocked_graph, session_id): |
| resp = api_client.post("/api/query/stream", json={ |
| "user_query": "top products", |
| "connector_id": "neon:public", |
| "session_id": session_id, |
| }) |
| raw = resp.text |
| events = [ |
| json.loads(line[len("data: "):]) |
| for line in raw.split("\n") |
| if line.startswith("data: ") |
| ] |
| token_events = [e for e in events if "token" in e] |
| assert len(token_events) > 0 |
|
|
