shudao-main/shudao-chat-py/tests/test_report_compat_proxy_body.py

import json
import importlib.util
import unittest
from pathlib import Path
from types import SimpleNamespace

from models.report import ReportCompleteFlowRequest


REPORT_COMPAT_PATH = Path(__file__).resolve().parents[1] / "routers" / "report_compat.py"
spec = importlib.util.spec_from_file_location("report_compat_under_test", REPORT_COMPAT_PATH)
report_compat = importlib.util.module_from_spec(spec)
spec.loader.exec_module(report_compat)

_build_aichat_complete_flow_body = report_compat._build_aichat_complete_flow_body


def _fake_request(user_id: int = 70430):
    return SimpleNamespace(
        state=SimpleNamespace(
            user=SimpleNamespace(user_id=user_id)
        )
    )


class ReportCompatProxyBodyTest(unittest.TestCase):
    def test_new_conversation_zero_is_forwarded_to_aichat(self):
        request_data = ReportCompleteFlowRequest(
            user_question="专业问题",
            ai_conversation_id=0,
        )

        body = _build_aichat_complete_flow_body(request_data, _fake_request())
        payload = json.loads(body.decode("utf-8"))

        self.assertEqual(payload["ai_conversation_id"], 0)

    def test_missing_conversation_id_is_forwarded_as_zero(self):
        request_data = ReportCompleteFlowRequest(
            user_question="专业问题",
        )

        body = _build_aichat_complete_flow_body(request_data, _fake_request())
        payload = json.loads(body.decode("utf-8"))

        self.assertEqual(payload["ai_conversation_id"], 0)

    def test_existing_conversation_id_is_preserved(self):
        request_data = ReportCompleteFlowRequest(
            user_question="继续追问",
            ai_conversation_id=11226,
        )

        body = _build_aichat_complete_flow_body(request_data, _fake_request())
        payload = json.loads(body.decode("utf-8"))

        self.assertEqual(payload["ai_conversation_id"], 11226)


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