more_dots/tests/test_stream_event_mode.py

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"