ai-contract-platform/backend/tests/test_chat_api.py

import json
import unittest

from fastapi import HTTPException
from fastapi.testclient import TestClient

from app.api.chat import (
    get_chat_legal_article_service,
    get_chat_prompt_runtime_service,
    get_chat_review_service,
    get_shared_ai_platform_client,
    get_tenant_service,
)
from app.api.conversations import get_shared_ai_platform_client as get_conversations_shared_ai_platform_client
from app.api.reviews import require_current_user
from app.main import app


class StubChatService:
    def __init__(self):
        self.stream_calls = []

    def stream_answer(self, current_user, conversation_id, message, review_context=None, template_key=None):
        self.stream_calls.append((current_user, conversation_id, message, review_context, template_key))
        if not message.strip():
            yield 'data: {"type":"error","message":"message 不能为空"}\n\n'
            return
        active_id = conversation_id or "conv-123"
        yield 'data: {"type":"text","content":"第一段"}\n\n'
        yield 'data: {"type":"text","content":"重点摘要：\\n- 关注违约金条款"}\n\n'
        yield 'data: {"type":"sources","sources":[{"title":"民法典第585条","content":"...","score":1.2}]}\n\n'
        yield f'data: {{"type":"end","conversation_id":"{active_id}"}}\n\n'

    def list_conversations(self, current_user, page, page_size):
        return {
            "items": [
                {
                    "id": "conv-123",
                    "title": "违约金条款咨询",
                    "message_count": 2,
                    "last_message_at": "2026-04-15T11:00:00+00:00",
                }
            ],
            "page": page,
            "page_size": page_size,
            "total": 1,
        }

    def get_conversation_detail(self, current_user, conversation_id):
        if conversation_id == "missing":
            raise HTTPException(status_code=404, detail="conversation not found")
        return {
            "id": conversation_id,
            "title": "违约金条款咨询",
            "messages": [
                {
                    "id": "msg-1",
                    "role": "user",
                    "content": "这个违约金条款有没有问题？",
                    "sources": [],
                    "created_at": "2026-04-15T11:00:00+00:00",
                },
                {
                    "id": "msg-2",
                    "role": "assistant",
                    "content": "存在偏高风险。",
                    "sources": [{"title": "民法典第585条", "content": "..."}],
                    "created_at": "2026-04-15T11:00:03+00:00",
                },
            ],
        }


class StubReviewService:
    def get_review_result(self, current_user, review_id):
        return {
            "risk_items": [
                {
                    "title": "违约金过高",
                    "legal_article_refs": [
                        {
                            "id": "lart-civil-585",
                            "source_id": "LAW-CIVIL-CODE-2020",
                            "source_title": "中华人民共和国民法典",
                            "article_no": "第五百八十五条",
                            "article_excerpt": "当事人可以约定一方违约时应当根据违约情况支付违约金。",
                            "authority": "全国人民代表大会",
                            "knowledge_base_id": "kb-phase14",
                            "knowledge_document_id": "doc-lart-civil-585",
                            "detail_path": "/legal-articles?article_id=lart-civil-585",
                            "match_reason": "命中风险关键词：违约金",
                            "tags": ["违约金"],
                        }
                    ],
                }
            ]
        }


class StubLegalArticleService:
    def __init__(self, sources=None):
        self.sources = sources or []
        self.calls = []

    def search_runtime_sources(self, current_user, query, limit=5):
        self.calls.append((current_user, query, limit))
        return self.sources[:limit]


class StubTenantService:
    def __init__(self, blocked=False):
        self.blocked = blocked
        self.calls = []

    def assert_billing_allows(self, current_user, action, requested_file_size_bytes=0):
        self.calls.append((action, requested_file_size_bytes))
        if self.blocked:
            raise HTTPException(status_code=402, detail="试用已到期，请升级套餐后继续使用 Chat 法律问答。")


class StubPromptRuntimeService:
    def __init__(self):
        self.record_calls = []

    def resolve_published_prompt(self, current_user, prompt_type, *, preferred_keys=None, fallback_prompt_version=None):
        return {
            "available": True,
            "prompt_asset_id": "prompt-001",
            "prompt_asset_version_id": "promptver-001",
            "prompt_key": "legal_chat_default",
            "prompt_type": prompt_type,
            "version_label": "v1.0",
            "prompt_version": "legal_chat_default.v1.0",
            "template_key": "legal_chat_default",
            "source": "published_prompt_asset",
        }

    def to_trace(self, prompt_runtime):
        return {
            key: prompt_runtime.get(key)
            for key in [
                "available",
                "prompt_asset_id",
                "prompt_asset_version_id",
                "prompt_key",
                "prompt_type",
                "version_label",
                "prompt_version",
                "template_key",
                "source",
            ]
        }

    def record_run_link(self, current_user, prompt_runtime, **kwargs):
        self.record_calls.append((current_user, prompt_runtime, kwargs))
        return {"id": "promptrun-001"}


class ChatApiTestCase(unittest.TestCase):
    def setUp(self):
        self.chat_service = StubChatService()
        self.tenant_service = StubTenantService()
        app.dependency_overrides[get_shared_ai_platform_client] = lambda: self.chat_service
        app.dependency_overrides[get_conversations_shared_ai_platform_client] = lambda: self.chat_service
        app.dependency_overrides[get_chat_review_service] = lambda: None
        app.dependency_overrides[get_chat_legal_article_service] = lambda: StubLegalArticleService()
        app.dependency_overrides[get_chat_prompt_runtime_service] = lambda: None
        app.dependency_overrides[get_tenant_service] = lambda: self.tenant_service
        self.client = TestClient(app)

    def tearDown(self):
        app.dependency_overrides.clear()

    def test_stream_chat_requires_authentication(self):
        response = self.client.post("/api/v1/chat/stream", json={"message": "违约金条款有没有问题"})

        self.assertEqual(response.status_code, 401)

    def test_stream_chat_returns_sse_events(self):
        app.dependency_overrides[require_current_user] = lambda: {"id": "user-1", "tenant_id": "tenant-1"}

        response = self.client.post("/api/v1/chat/stream", json={"message": "违约金条款有没有问题"})

        self.assertEqual(response.status_code, 200)
        self.assertIn("text/event-stream", response.headers["content-type"])

        events = [line.removeprefix("data: ") for line in response.text.split("\n\n") if line.startswith("data: ")]
        payloads = [json.loads(item) for item in events]

        self.assertEqual(payloads[0]["type"], "text")
        self.assertEqual(payloads[1]["type"], "text")
        self.assertEqual(payloads[2]["type"], "sources")
        self.assertEqual(payloads[2]["sources"][0]["score"], 1.2)
        self.assertEqual(payloads[3]["type"], "end")
        self.assertEqual(payloads[3]["conversation_id"], "conv-123")
        self.assertEqual(self.chat_service.stream_calls[0][0]["id"], "user-1")
        self.assertEqual(self.tenant_service.calls[0][0], "chat")

    def test_stream_chat_returns_billing_block_message(self):
        app.dependency_overrides[require_current_user] = lambda: {"id": "user-1", "tenant_id": "tenant-1"}
        app.dependency_overrides[get_tenant_service] = lambda: StubTenantService(blocked=True)

        response = self.client.post("/api/v1/chat/stream", json={"message": "违约金条款有没有问题"})

        self.assertEqual(response.status_code, 402)
        self.assertIn("试用已到期", response.json()["message"])

    def test_stream_chat_accepts_review_context(self):
        app.dependency_overrides[require_current_user] = lambda: {"id": "user-1", "tenant_id": "tenant-1"}

        response = self.client.post(
            "/api/v1/chat/stream",
            json={
                "message": "请继续解释这个风险项",
                "review_context": {
                    "reviewId": "review-123",
                    "metadata": {
                        "reviewTitle": "采购合同审查",
                        "focusLabel": "自动续约",
                        "actionLabel": "解释风险",
                    },
                },
            },
        )

        self.assertEqual(response.status_code, 200)
        self.assertEqual(self.chat_service.stream_calls[0][3]["reviewId"], "review-123")
        self.assertEqual(self.chat_service.stream_calls[0][3]["metadata"]["focusLabel"], "自动续约")

    def test_stream_chat_passes_published_prompt_template_key(self):
        app.dependency_overrides[require_current_user] = lambda: {"id": "user-1", "tenant_id": "tenant-1"}
        prompt_runtime_service = StubPromptRuntimeService()
        app.dependency_overrides[get_chat_prompt_runtime_service] = lambda: prompt_runtime_service

        response = self.client.post("/api/v1/chat/stream", json={"message": "check clause"})

        self.assertEqual(response.status_code, 200)
        review_context = self.chat_service.stream_calls[0][3]
        self.assertEqual(self.chat_service.stream_calls[0][4], "legal_chat_default")
        self.assertEqual(review_context["promptRuntime"]["prompt_version"], "legal_chat_default.v1.0")
        self.assertEqual(prompt_runtime_service.record_calls[0][2]["run_type"], "legal_chat")

    def test_stream_chat_adds_review_legal_article_sources(self):
        app.dependency_overrides[require_current_user] = lambda: {"id": "user-1", "tenant_id": "tenant-1"}
        app.dependency_overrides[get_chat_review_service] = lambda: StubReviewService()

        response = self.client.post(
            "/api/v1/chat/stream",
            json={
                "message": "请解释违约金风险",
                "review_context": {
                    "reviewId": "review-123",
                    "metadata": {"focusLabel": "违约金"},
                },
            },
        )

        self.assertEqual(response.status_code, 200)
        events = [line.removeprefix("data: ") for line in response.text.split("\n\n") if line.startswith("data: ")]
        payloads = [json.loads(item) for item in events]
        sources_event = next(item for item in payloads if item["type"] == "sources")
        legal_source = next(item for item in sources_event["sources"] if item.get("source_type") == "legal_article")

        self.assertEqual(legal_source["source_id"], "lart-civil-585")
        self.assertEqual(legal_source["detail_path"], "/legal-articles?article_id=lart-civil-585")
        self.assertEqual(self.chat_service.stream_calls[0][3]["legalArticleRefs"][0]["id"], "lart-civil-585")

    def test_stream_chat_adds_runtime_legal_article_sources(self):
        app.dependency_overrides[require_current_user] = lambda: {"id": "user-1", "tenant_id": "tenant-1"}
        legal_article_service = StubLegalArticleService(
            [
                {
                    "title": "中华人民共和国网络安全法 第一条",
                    "content": "为了保障网络安全，维护网络空间主权和国家安全、社会公共利益，制定本法。",
                    "score": 12.5,
                    "source_type": "legal_article",
                    "source_id": "lart-cyber-001",
                    "detail_path": "/legal-articles?article_id=lart-cyber-001",
                    "article_no": "第一条",
                    "authority": "全国人民代表大会常务委员会",
                    "metadata": {
                        "source_id": "FLK-CYBER-SECURITY",
                        "source_title": "中华人民共和国网络安全法",
                        "match_reason": "命中条号；命中法源",
                        "tags": ["网络安全"],
                    },
                }
            ]
        )
        app.dependency_overrides[get_chat_legal_article_service] = lambda: legal_article_service

        response = self.client.post("/api/v1/chat/stream", json={"message": "网络安全法第一条是什么"})

        self.assertEqual(response.status_code, 200)
        events = [line.removeprefix("data: ") for line in response.text.split("\n\n") if line.startswith("data: ")]
        payloads = [json.loads(item) for item in events]
        sources_event = next(item for item in payloads if item["type"] == "sources")
        legal_source = next(item for item in sources_event["sources"] if item.get("source_type") == "legal_article")

        self.assertEqual(legal_source["source_id"], "lart-cyber-001")
        self.assertEqual(legal_source["detail_path"], "/legal-articles?article_id=lart-cyber-001")
        self.assertEqual(legal_article_service.calls[0][1], "网络安全法第一条是什么")
        review_context = self.chat_service.stream_calls[0][3]
        self.assertEqual(review_context["legalArticleSearchSources"][0]["source_id"], "lart-cyber-001")
        self.assertEqual(review_context["metadata"]["legalArticleSearchSourceCount"], 1)

    def test_list_conversations_returns_paginated_payload(self):
        app.dependency_overrides[require_current_user] = lambda: {"id": "user-1", "tenant_id": "tenant-1"}

        response = self.client.get("/api/v1/conversations?page=1&page_size=20")

        self.assertEqual(response.status_code, 200)
        payload = response.json()["data"]
        self.assertEqual(payload["total"], 1)
        self.assertEqual(payload["items"][0]["id"], "conv-123")

    def test_get_conversation_detail_returns_messages(self):
        app.dependency_overrides[require_current_user] = lambda: {"id": "user-1", "tenant_id": "tenant-1"}

        response = self.client.get("/api/v1/conversations/conv-123")

        self.assertEqual(response.status_code, 200)
        payload = response.json()["data"]
        self.assertEqual(payload["id"], "conv-123")
        self.assertEqual(len(payload["messages"]), 2)
        self.assertEqual(payload["messages"][1]["sources"][0]["title"], "民法典第585条")

    def test_get_conversation_detail_returns_404_when_missing(self):
        app.dependency_overrides[require_current_user] = lambda: {"id": "user-1", "tenant_id": "tenant-1"}

        response = self.client.get("/api/v1/conversations/missing")

        self.assertEqual(response.status_code, 404)


if __name__ == "__main__":
    unittest.main()