more_dots/tests/test_conversation_handling.py

from fastapi import FastAPI
from fastapi.testclient import TestClient

from api.dependencies import get_workflow_manager
from api.endpoints import router
from workflows.workflow_manager import WorkflowType


class StubWorkflowManager:
    def __init__(self):
        self.calls = []

    def execute_workflow(self, workflow_type, user_input, session_id=None, **kwargs):
        self.calls.append(
            {
                "workflow_type": workflow_type,
                "user_input": user_input,
                "session_id": session_id,
                "kwargs": kwargs,
            }
        )
        return {
            "session_id": session_id,
            "workflow_type": workflow_type.value if isinstance(workflow_type, WorkflowType) else str(workflow_type),
            "result": {"ok": True, "context": {"final_sql": "SELECT 1"}},
        }


class StubMessageStorage:
    def __init__(self):
        self.enabled = True
        self.created = []
        self.updated = []
        self.saved = []
        self.existing = {}
        self.create_should_fail = False
        self.update_should_fail = False

    def create_conversation(self, conversation_id, user, name, status, introduction, created_at, updated_at):
        if self.create_should_fail:
            return False
        self.created.append(
            {
                "conversation_id": conversation_id,
                "user": user,
                "name": name,
                "status": status,
                "introduction": introduction,
                "created_at": created_at,
                "updated_at": updated_at,
            }
        )
        self.existing[conversation_id] = self.created[-1]
        return True

    def get_conversation_by_id(self, conversation_id):
        return self.existing.get(conversation_id)

    def update_conversation_updated_at(self, conversation_id, updated_at):
        self.updated.append({"conversation_id": conversation_id, "updated_at": updated_at})
        if self.update_should_fail:
            return False
        return conversation_id in self.existing

    def save_message(self, **kwargs):
        self.saved.append(kwargs)
        return True


def _build_client(workflow_manager, monkeypatch, storage):
    app = FastAPI()
    app.include_router(router)
    app.dependency_overrides[get_workflow_manager] = lambda: workflow_manager
    monkeypatch.setattr("api.endpoints.get_message_storage", lambda: storage)
    return TestClient(app)


def test_run_workflow_creates_conversation_when_missing_id(monkeypatch):
    workflow_manager = StubWorkflowManager()
    storage = StubMessageStorage()
    client = _build_client(workflow_manager, monkeypatch, storage)

    response = client.post(
        "/api/workflows",
        json={
            "query": "查询订单",
            "inputs": {},
            "response_mode": "blocking",
            "user": "tester",
            "conversation_id": None,
            "files": [],
        },
    )

    assert response.status_code == 200
    assert len(storage.created) == 1
    created = storage.created[0]
    assert created["status"] == "normal"
    assert created["name"] == "查询订单"
    assert created["introduction"] is None
    assert created["created_at"] == created["updated_at"]
    assert workflow_manager.calls[0]["session_id"] == created["conversation_id"]
    assert response.json()["session_id"] == created["conversation_id"]
    assert storage.saved[0]["created_at"] == created["created_at"]
    assert storage.saved[0]["updated_at"] == created["updated_at"]


def test_run_workflow_first_turn_name_uses_first_20_chars(monkeypatch):
    workflow_manager = StubWorkflowManager()
    storage = StubMessageStorage()
    client = _build_client(workflow_manager, monkeypatch, storage)

    response = client.post(
        "/api/workflows",
        json={
            "query": "12345678901234567890EXTRA_TEXT",
            "inputs": {},
            "response_mode": "blocking",
            "user": "tester",
            "conversation_id": None,
            "files": [],
        },
    )

    assert response.status_code == 200
    assert storage.created[0]["name"] == "12345678901234567890"


def test_run_workflow_updates_existing_conversation(monkeypatch):
    workflow_manager = StubWorkflowManager()
    storage = StubMessageStorage()
    storage.existing["cid-exists"] = {
        "conversation_id": "cid-exists",
        "user": "tester",
        "name": "old",
        "status": "normal",
        "created_at": 1,
        "updated_at": 1,
    }
    client = _build_client(workflow_manager, monkeypatch, storage)

    response = client.post(
        "/api/workflows",
        json={
            "query": "查询订单",
            "inputs": {},
            "response_mode": "blocking",
            "user": "tester",
            "conversation_id": "cid-exists",
            "files": [],
        },
    )

    assert response.status_code == 200
    assert storage.updated and storage.updated[0]["conversation_id"] == "cid-exists"
    assert workflow_manager.calls[0]["session_id"] == "cid-exists"


def test_run_workflow_returns_400_when_provided_conversation_id_not_found(monkeypatch):
    workflow_manager = StubWorkflowManager()
    storage = StubMessageStorage()
    client = _build_client(workflow_manager, monkeypatch, storage)

    response = client.post(
        "/api/workflows",
        json={
            "query": "查询订单",
            "inputs": {},
            "response_mode": "blocking",
            "user": "tester",
            "conversation_id": "cid-missing",
            "files": [],
        },
    )

    assert response.status_code == 400
    assert response.json()["detail"]["code"] == "CONVERSATION_NOT_FOUND"


def test_run_workflow_returns_500_when_conversation_create_fails(monkeypatch):
    workflow_manager = StubWorkflowManager()
    storage = StubMessageStorage()
    storage.create_should_fail = True
    client = _build_client(workflow_manager, monkeypatch, storage)

    response = client.post(
        "/api/workflows",
        json={
            "query": "查询订单",
            "inputs": {},
            "response_mode": "blocking",
            "user": "tester",
            "conversation_id": None,
            "files": [],
        },
    )

    assert response.status_code == 500
    assert response.json()["detail"]["code"] == "CONVERSATION_CREATE_FAILED"


def test_run_workflow_returns_500_when_conversation_update_fails(monkeypatch):
    workflow_manager = StubWorkflowManager()
    storage = StubMessageStorage()
    storage.existing["cid-exists"] = {
        "conversation_id": "cid-exists",
        "user": "tester",
        "name": "old",
        "status": "normal",
        "created_at": 1,
        "updated_at": 1,
    }
    storage.update_should_fail = True
    client = _build_client(workflow_manager, monkeypatch, storage)

    response = client.post(
        "/api/workflows",
        json={
            "query": "查询订单",
            "inputs": {},
            "response_mode": "blocking",
            "user": "tester",
            "conversation_id": "cid-exists",
            "files": [],
        },
    )

    assert response.status_code == 500
    assert response.json()["detail"]["code"] == "CONVERSATION_UPDATE_FAILED"