init

tests/test_stream_event_mode.py

@@ -0,0 +1,154 @@
+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 "<strong>Question:</strong> test" in dto.answer
+    assert "<table>" in dto.answer
+    assert "<th>value</th>" in dto.answer
+    assert "<td>1</td>" in dto.answer
+    assert "<strong>Rows:</strong> 1" in dto.answer
+    assert "<strong>Data Version:</strong> 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 "<div" not in dto.answer
+    assert "<strong>" not in dto.answer
+    end_dto = ChatMessageResponseDTO.model_validate_json(lines[-1][len("data: "):])
+    assert end_dto.event == "message_end"
+    assert storage.saved[0]["answer"] == dto.answer
+    assert any(item.startswith("stream.sql_executed") for item in storage.saved[0]["logs"])
+
+
+def test_stream_error_still_returns_plain_text(monkeypatch):
+    monkeypatch.setattr("api.endpoints.SrApiQueryTool", FakeSrApiTool)
+    monkeypatch.setattr("api.endpoints.get_message_storage", lambda: DisabledMessageStorage())
+    client = _build_client(StubWorkflowManager())
+
+    with client.stream("POST", "/api/workflows/stream", json={**_payload(), "query": "   "}) as response:
+        payload = json.loads(response.read().decode("utf-8"))
+
+    assert response.status_code == 400
+    assert payload["detail"]["code"] == "INVALID_REQUEST"
+