import json from fastapi import FastAPI from fastapi.testclient import TestClient from api.dependencies import get_workflow_manager from api.endpoints import router from schemas.chat_message_response import ChatMessageResponseDTO from workflows.workflow_manager import WorkflowType class StubWorkflowManager: def execute_workflow(self, workflow_type, user_input, session_id=None, **kwargs): return { "session_id": session_id or "cid-1", "workflow_type": workflow_type.value if isinstance(workflow_type, WorkflowType) else str(workflow_type), "result": {"context": {"final_sql": "SELECT 1"}}, } class FakeSrApiTool: def run(self, payload): return '{"total": 1, "data": [{"value": 1, "etl_time": "2026-03-20 11:27:53"}]}' class FakeEmptySrApiTool: def run(self, payload): return '{"total": 0, "data": []}' class DisabledMessageStorage: enabled = False def save_message(self, **kwargs): return False class CaptureMessageStorage: enabled = True def __init__(self): self.saved = [] self.existing = { "cid-test-1": { "conversation_id": "cid-test-1", "user": "tester", "name": "test", "status": "normal", "created_at": 1, "updated_at": 1, } } def get_conversation_by_id(self, conversation_id): return self.existing.get(conversation_id) def update_conversation_updated_at(self, conversation_id, updated_at): if conversation_id not in self.existing: return False self.existing[conversation_id]["updated_at"] = updated_at return True def save_message(self, **kwargs): self.saved.append(kwargs) return True def _build_client(workflow_manager) -> TestClient: app = FastAPI() app.include_router(router) app.dependency_overrides[get_workflow_manager] = lambda: workflow_manager return TestClient(app) def _payload(): return { "query": "test", "inputs": {}, "response_mode": "streaming", "user": "tester", "conversation_id": "cid-test-1", "files": [], } def test_stream_returns_chat_message_response_dto_chunks(monkeypatch): monkeypatch.setattr("api.endpoints.SrApiQueryTool", FakeSrApiTool) storage = CaptureMessageStorage() monkeypatch.setattr("api.endpoints.get_message_storage", lambda: storage) client = _build_client(StubWorkflowManager()) with client.stream("POST", "/api/workflows/stream", json=_payload()) as response: body = response.read().decode("utf-8") assert response.status_code == 200 lines = [line for line in body.splitlines() if line.strip()] assert all(not line.startswith("event:") for line in lines) assert lines[0].startswith("data: ") dto = ChatMessageResponseDTO.model_validate_json(lines[0][len("data: "):]) assert dto.event == "message" assert dto.conversation_id == "cid-test-1" assert "Question: test" in dto.answer assert "
| value | " in dto.answer assert "1 | " in dto.answer assert "Rows: 1" in dto.answer assert "Data Version: 2026-03-20 11:27:53" in dto.answer assert dto.task_id assert dto.message_id end_dto = ChatMessageResponseDTO.model_validate_json(lines[-1][len("data: "):]) assert end_dto.event == "message_end" assert end_dto.answer == "" assert len(storage.saved) == 1 assert storage.saved[0]["answer"] == dto.answer assert storage.saved[0]["logs"][0].startswith("stream.start") assert any(item.startswith("stream.sql_executed") for item in storage.saved[0]["logs"]) def test_stream_returns_no_data_when_sql_data_is_empty(monkeypatch): monkeypatch.setattr("api.endpoints.SrApiQueryTool", FakeEmptySrApiTool) storage = CaptureMessageStorage() monkeypatch.setattr("api.endpoints.get_message_storage", lambda: storage) client = _build_client(StubWorkflowManager()) with client.stream("POST", "/api/workflows/stream", json=_payload()) as response: body = response.read().decode("utf-8") assert response.status_code == 200 lines = [line for line in body.splitlines() if line.strip()] dto = ChatMessageResponseDTO.model_validate_json(lines[0][len("data: "):]) assert "Question: test" in dto.answer assert "No data~" in dto.answer assert "Rows: 0" in dto.answer assert "Data Version: Unknown" in dto.answer assert "
|---|