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

import json
import importlib.util
from pathlib import Path
import unittest
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)


class ReportCompatProxyBodyTests(unittest.TestCase):
    def _request(self, user_id=70430):
        return SimpleNamespace(state=SimpleNamespace(user=SimpleNamespace(user_id=user_id)))

    def _payload_for(self, ai_conversation_id, request_user_id=70430, payload_user_id=None):
        request_data = ReportCompleteFlowRequest(
            user_question="history save regression",
            ai_conversation_id=ai_conversation_id,
            user_id=payload_user_id,
        )
        body = report_compat._build_aichat_complete_flow_body(
            request_data,
            self._request(user_id=request_user_id),
        )
        return json.loads(body.decode("utf-8"))

    def test_new_conversation_is_forwarded_as_zero(self):
        payload = self._payload_for(0)

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

    def test_missing_conversation_id_is_forwarded_as_zero(self):
        payload = self._payload_for(None)

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

    def test_existing_conversation_id_is_preserved(self):
        payload = self._payload_for(12345)

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

    def test_request_user_id_is_forwarded_to_aichat(self):
        payload = self._payload_for(0, request_user_id=70430)

        self.assertEqual(payload["user_id"], 70430)

    def test_payload_user_id_takes_precedence(self):
        payload = self._payload_for(0, request_user_id=70430, payload_user_id=88)

        self.assertEqual(payload["user_id"], 88)


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