xiaobai-review/tests/test_llm_transport.py

from __future__ import annotations

import io
import json
import unittest
import urllib.error
from pathlib import Path
from unittest.mock import patch

from backend.features.heaven.agent import HeavenAgentError, interpret_heaven
from backend.features.mentor.agent import MentorAgentError, MentorSkill, stream_with_mentor
from backend.features.review.agent import ReviewAssistantError, stream_review_assistant
from backend.features.screener.compiler import LLMCompilerError, test_llm_connection
from backend.llm import transport


ROOT = Path(__file__).resolve().parents[1]


class FakeResponse:
    def __init__(self, *, payload: bytes = b"", lines: list[bytes] | None = None) -> None:
        self.payload = payload
        self.lines = lines or []

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        return False

    def read(self) -> bytes:
        return self.payload

    def __iter__(self):
        return iter(self.lines)


class OpenAITransportTests(unittest.TestCase):
    def test_chat_completion_builds_one_openai_compatible_request(self) -> None:
        captured = {}
        response = FakeResponse(
            payload=json.dumps(
                {"choices": [{"message": {"content": "OK"}}]}
            ).encode("utf-8")
        )

        def open_request(request, timeout):
            captured["url"] = request.full_url
            captured["headers"] = request.headers
            captured["payload"] = json.loads(request.data.decode("utf-8"))
            captured["timeout"] = timeout
            return response

        with patch("backend.llm.transport.urllib.request.urlopen", side_effect=open_request):
            result = transport.chat_completion(
                api_key="secret",
                base_url="https://example.test/v1/",
                model="model",
                messages=[{"role": "user", "content": "ping"}],
                timeout=17,
                user_agent="XiaobaiReviewWeb/test",
            )

        self.assertEqual(result.content, "OK")
        self.assertGreaterEqual(result.latency_ms, 0)
        self.assertEqual(captured["url"], "https://example.test/v1/chat/completions")
        self.assertEqual(captured["payload"]["stream"], False)
        self.assertEqual(captured["headers"]["Authorization"], "Bearer secret")
        self.assertEqual(captured["timeout"], 17)

    def test_stream_completion_parses_deltas_and_ignores_final_snapshot(self) -> None:
        response = FakeResponse(
            lines=[
                b'data: {"choices":[{"delta":{"content":"first"}}]}\n',
                b'data: {"choices":[{"delta":{"content":" second"}}]}\n',
                b'data: {"choices":[{"message":{"content":"first second"}}]}\n',
                b"data: [DONE]\n",
            ]
        )
        with patch("backend.llm.transport.urllib.request.urlopen", return_value=response):
            chunks = list(
                transport.stream_chat_completion(
                    api_key="secret",
                    base_url="https://example.test/v1",
                    model="model",
                    messages=[],
                    timeout=17,
                    user_agent="XiaobaiReviewWeb/test",
                )
            )
        self.assertEqual(chunks, ["first", " second"])

    def test_empty_stream_has_a_stable_transport_error(self) -> None:
        response = FakeResponse(lines=[b"data: [DONE]\n"])
        with patch("backend.llm.transport.urllib.request.urlopen", return_value=response):
            with self.assertRaises(transport.OpenAIEmptyResponseError):
                list(
                    transport.stream_chat_completion(
                        api_key="secret",
                        base_url="https://example.test/v1",
                        model="model",
                        messages=[],
                        timeout=17,
                        user_agent="XiaobaiReviewWeb/test",
                    )
                )

    def test_http_error_keeps_code_and_sanitized_provider_detail(self) -> None:
        error = urllib.error.HTTPError(
            "https://example.test/v1/chat/completions",
            429,
            "rate limited",
            {},
            io.BytesIO(b'{"error":{"message":"capacity"}}'),
        )
        self.addCleanup(error.close)
        with patch("backend.llm.transport.urllib.request.urlopen", side_effect=error):
            with self.assertRaises(transport.OpenAIHTTPError) as caught:
                transport.chat_completion(
                    api_key="secret",
                    base_url="https://example.test/v1",
                    model="model",
                    messages=[],
                    timeout=17,
                    user_agent="XiaobaiReviewWeb/test",
                )
        self.assertEqual(caught.exception.code, 429)
        self.assertEqual(
            caught.exception.describe("模型调用失败"),
            "模型调用失败（HTTP 429）：capacity",
        )

    def test_feature_agents_have_no_direct_provider_transport(self) -> None:
        paths = (
            "backend/features/mentor/agent.py",
            "backend/features/heaven/agent.py",
            "backend/features/review/agent.py",
            "backend/features/screener/compiler.py",
        )
        for relative in paths:
            source = (ROOT / relative).read_text(encoding="utf-8")
            with self.subTest(path=relative):
                self.assertNotIn("urllib.request", source)
                self.assertNotIn("/chat/completions", source)
                self.assertIn("llm_transport.", source)


class FeatureErrorMappingTests(unittest.TestCase):
    def test_feature_specific_http_messages_are_preserved(self) -> None:
        error = transport.OpenAIHTTPError(429, "capacity")
        skill = MentorSkill(
            skill_id="test",
            name="测试老师",
            description="",
            tagline="",
            focus=(),
            content="",
            path=Path("SKILL.md"),
        )
        with patch(
            "backend.features.mentor.agent.llm_transport.stream_chat_completion",
            side_effect=error,
        ):
            with self.assertRaisesRegex(
                MentorAgentError, "问师模型调用失败（HTTP 429）：capacity"
            ):
                list(stream_with_mentor(skill, {}, "问题", [], "key", "https://x", "m"))
        with patch(
            "backend.features.heaven.agent.llm_transport.chat_completion",
            side_effect=error,
        ):
            with self.assertRaisesRegex(
                HeavenAgentError, "问天模型调用失败（HTTP 429）：capacity"
            ):
                interpret_heaven("heart", {}, "key", "https://x", "m")
        with patch(
            "backend.features.review.agent.llm_transport.stream_chat_completion",
            side_effect=error,
        ):
            with self.assertRaisesRegex(
                ReviewAssistantError, "智能解读服务暂不可用（429）"
            ):
                list(stream_review_assistant({}, "问题", [], "key", "https://x", "m"))
        with patch(
            "backend.features.screener.compiler.llm_transport.chat_completion",
            side_effect=error,
        ):
            with self.assertRaisesRegex(
                LLMCompilerError, "模型连接测试失败（HTTP 429）：capacity"
            ):
                test_llm_connection("key", "https://x", "m")


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