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({}), "")

    def test_extract_titles_comma_separated(self):
        """§19.23: 逗号分隔（全角/半角）→ 多条，逐段去空白。"""
        from services.ai_text_service import extract_titles_from_response

        full = {"choices": [{"message": {"content": "甲款，乙款，丙款，丁款"}}]}
        self.assertEqual(extract_titles_from_response(full),
                         ["甲款", "乙款", "丙款", "丁款"])
        half = {"choices": [{"message": {"content": "甲款, 乙款 ,丙款"}}]}
        self.assertEqual(extract_titles_from_response(half), ["甲款", "乙款", "丙款"])

    def test_extract_titles_mixed_comma_newline_and_cleaning(self):
        """§19.23: 逗号 + 换行混用都拆开；逐段去序号/引号。"""
        from services.ai_text_service import extract_titles_from_response

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

    def test_extract_titles_markdown_table_not_exploded(self):
        """§19.23: 表头单行无逗号 → 至多 1 条，不再炸成表头/分隔/多行垃圾。"""
        from services.ai_text_service import extract_titles_from_response

        data = {"choices": [{"message": {"content": "| 編號 | 標題 | 字元數估算 |"}}]}
        self.assertEqual(len(extract_titles_from_response(data)), 1)

    def test_extract_titles_empty_when_no_text(self):
        from services.ai_text_service import extract_titles_from_response

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


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("标题")

    def test_generate_texts_returns_all_titles_text_only(self):
        """§17.1: 一次请求 → 多条；纯文本（payload 不含 image_url）。"""
        from services.ai_text_service import AiTextClient

        session = _FakeSession({"choices": [{"message": {"content":
            "标题一\n标题二\n标题三"}}]})
        client = AiTextClient(CHAT_CFG, session=session)

        titles = client.generate_texts("生成 3 条标题")

        self.assertEqual(titles, ["标题一", "标题二", "标题三"])
        content = session.posted["json"]["messages"][0]["content"]
        self.assertEqual([p["type"] for p in content], ["text"])   # 无图

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

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


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