cmbot/tests/test_ai_text_service.py

"""Tests for the AI text (title) service (docs/11 §17)."""
import sys
import unittest
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent / "src"))


CHAT_CFG = {"url": "https://relay.example.com/v1/chat/completions",
            "model": "gpt-4o", "api_key": "sk-x", "api_type": "chat"}
GEMINI_CFG = {"url": "https://relay.example.com/v1beta/models/{model}:generateContent",
              "model": "gemini-2.5-flash", "api_key": "sk-x", "api_type": "gemini"}
IMAGES_EDITS_CFG = {"url": "https://relay.example.com/v1/images/edits",
                    "model": "img", "api_key": "sk-x", "api_type": "images_edits"}


class TestExtractText(unittest.TestCase):
    def test_chat_string_content(self):
        from services.ai_text_service import extract_text_from_response

        data = {"choices": [{"message": {"content": "纯棉宽松短袖T恤"}}]}
        self.assertEqual(extract_text_from_response(data), "纯棉宽松短袖T恤")

    def test_chat_list_content(self):
        from services.ai_text_service import extract_text_from_response

        data = {"choices": [{"message": {"content": [
            {"type": "text", "text": "印花连衣裙"}]}}]}
        self.assertEqual(extract_text_from_response(data), "印花连衣裙")

    def test_gemini_parts(self):
        from services.ai_text_service import extract_text_from_response

        data = {"candidates": [{"content": {"parts": [{"text": "碎花雪纺衫"}]}}]}
        self.assertEqual(extract_text_from_response(data), "碎花雪纺衫")

    def test_multiline_takes_first_and_strips_numbering_quotes(self):
        from services.ai_text_service import extract_text_from_response

        data = {"choices": [{"message": {"content":
                "1. 「韩版宽松卫衣」\n2. 复古工装外套\n3. 简约百搭T恤"}}]}
        self.assertEqual(extract_text_from_response(data), "韩版宽松卫衣")

    def test_missing_text_returns_empty(self):
        from services.ai_text_service import extract_text_from_response

        self.assertEqual(extract_text_from_response({"choices": []}), "")
        self.assertEqual(extract_text_from_response({}), "")


class TestBuildTextPayload(unittest.TestCase):
    def test_chat_with_image_includes_image_url(self):
        from services.ai_text_service import build_text_payload

        payload = build_text_payload(CHAT_CFG, "写标题", "data:image/png;base64,AAA")
        content = payload["messages"][0]["content"]
        types = [part["type"] for part in content]
        self.assertEqual(types, ["text", "image_url"])
        self.assertEqual(content[1]["image_url"]["url"], "data:image/png;base64,AAA")

    def test_chat_without_image_text_only(self):
        from services.ai_text_service import build_text_payload

        payload = build_text_payload(CHAT_CFG, "写标题")
        content = payload["messages"][0]["content"]
        self.assertEqual([p["type"] for p in content], ["text"])

    def test_gemini_text_only_modality(self):
        from services.ai_text_service import build_text_payload

        payload = build_text_payload(GEMINI_CFG, "写标题",
                                     "data:image/jpeg;base64,BBB")
        self.assertEqual(
            payload["generationConfig"]["responseModalities"], ["TEXT"])
        parts = payload["contents"][0]["parts"]
        self.assertIn("inlineData", parts[1])

    def test_images_edits_rejected(self):
        from services.ai_text_service import AiTextServiceError, build_text_payload

        with self.assertRaises(AiTextServiceError):
            build_text_payload(IMAGES_EDITS_CFG, "写标题")


class _FakeResponse:
    def __init__(self, payload):
        self._payload = payload

    def raise_for_status(self):
        pass

    def json(self):
        return self._payload


class _FakeSession:
    def __init__(self, payload):
        self._payload = payload
        self.posted = None

    def post(self, url, headers=None, json=None, timeout=None):
        self.posted = {"url": url, "headers": headers, "json": json}
        return _FakeResponse(self._payload)


class TestGenerateText(unittest.TestCase):
    def test_generate_text_posts_and_returns_title(self):
        from services.ai_text_service import AiTextClient

        session = _FakeSession({"choices": [{"message": {"content": "夏季冰丝防晒衣"}}]})
        client = AiTextClient(CHAT_CFG, session=session)

        # No image_path -> text-only payload; still returns the parsed title.
        title = client.generate_text("给这件衣服起标题")

        self.assertEqual(title, "夏季冰丝防晒衣")
        self.assertIn("chat/completions", session.posted["url"])
        self.assertEqual(session.posted["headers"]["Authorization"], "Bearer sk-x")

    def test_generate_text_images_edits_raises(self):
        from services.ai_text_service import AiTextClient, AiTextServiceError

        client = AiTextClient(IMAGES_EDITS_CFG, session=_FakeSession({}))
        with self.assertRaises(AiTextServiceError):
            client.generate_text("标题")

    def test_generate_text_empty_response_raises(self):
        from services.ai_text_service import AiTextClient, AiTextServiceError

        client = AiTextClient(CHAT_CFG, session=_FakeSession({"choices": []}))
        with self.assertRaises(AiTextServiceError):
            client.generate_text("标题")


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