xiaoxia-saas/tests/unit/test_ai_client_image.py

"""#2170 Seedream 图片生成 + 方舟信任链单测。"""

from __future__ import annotations

from unittest.mock import MagicMock, patch

import httpx

from packages.shared.ai_client import DoubaoClient


def _make_client(**overrides):
    client = DoubaoClient.__new__(DoubaoClient)
    client.api_key = overrides.get("api_key", "test-key")
    client.base_url = overrides.get("base_url", "https://ark.cn-beijing.volces.com/api/v3")
    client.model = "doubao-model"
    client.vision_model = "doubao-vision"
    client.embedding_model = "doubao-embedding"
    client.image_model = overrides.get("image_model", "doubao-seedream-5-0-pro-260628")
    client.image_timeout = overrides.get("image_timeout", 120)
    client.timeout = overrides.get("timeout", 30)
    client.max_retries = overrides.get("max_retries", 0)
    client.last_video_error = {}
    client.last_image_error = {}
    return client


# ── Seedream 图片生成单测 ──────────────────────────────────────────


class TestImageGenerationHappyPath:
    def test_returns_none_when_no_api_key(self):
        client = _make_client(api_key="")
        assert client.image_generation("p") is None
        err = client.get_last_image_error()
        assert err["error_code"] == "auth_error"

    def test_returns_none_on_empty_prompt(self):
        client = _make_client()
        assert client.image_generation("   ") is None
        err = client.get_last_image_error()
        assert err["error_code"] == "invalid_param"

    def test_text_to_image_success(self):
        client = _make_client()
        captured = {}
        ok_resp = MagicMock()
        ok_resp.status_code = 200
        ok_resp.json.return_value = {"data": [{"url": "https://cdn.example.com/i.png"}], "usage": {"tokens": 1}}
        ok_resp.raise_for_status = MagicMock()
        ok_resp.text = ""

        def fake_post(url, **kwargs):
            captured["url"] = url
            captured["json"] = kwargs.get("json")
            return ok_resp

        with patch("packages.shared.ai_client.httpx.post", side_effect=fake_post):
            result = client.image_generation("一只可爱的猫", size="1K")
        assert result is not None
        assert result["url"] == "https://cdn.example.com/i.png"
        assert "/images/generations" in captured["url"]
        assert captured["json"]["model"] == "doubao-seedream-5-0-pro-260628"
        assert captured["json"]["size"] == "1K"
        assert "image" not in captured["json"]

    def test_image_to_image_single_ref_passed_as_string(self):
        client = _make_client()
        captured = {}
        ok_resp = MagicMock(status_code=200)
        ok_resp.json.return_value = {"data": [{"url": "https://cdn.example.com/out.png"}]}
        ok_resp.raise_for_status = MagicMock()
        ok_resp.text = ""

        def fake_post(url, **kwargs):
            captured["json"] = kwargs.get("json")
            return ok_resp

        with patch("packages.shared.ai_client.httpx.post", side_effect=fake_post):
            client.image_generation("保持五官", reference_images=["https://img/x.jpg"])
        assert captured["json"]["image"] == "https://img/x.jpg"

    def test_image_to_image_multiple_refs_passed_as_list(self):
        client = _make_client()
        captured = {}
        ok_resp = MagicMock(status_code=200)
        ok_resp.json.return_value = {"data": [{"url": "https://cdn.example.com/out.png"}]}
        ok_resp.raise_for_status = MagicMock()

        def fake_post(url, **kwargs):
            captured["json"] = kwargs.get("json")
            return ok_resp

        refs = [f"https://img/{i}.jpg" for i in range(3)]
        with patch("packages.shared.ai_client.httpx.post", side_effect=fake_post):
            client.image_generation("保持", reference_images=refs)
        assert captured["json"]["image"] == refs

    def test_400_sensitive_returns_portrait_intercept(self):
        client = _make_client(max_retries=0)
        bad_resp = MagicMock(status_code=400)
        bad_resp.text = '{"error":{"code":"ContentRisk","message":"sensitive content detected"}}'
        bad_resp.json.return_value = {"error": {"code": "ContentRisk"}}
        bad_resp.raise_for_status.side_effect = httpx.HTTPStatusError("bad", request=MagicMock(), response=bad_resp)
        with patch("packages.shared.ai_client.httpx.post", return_value=bad_resp):
            assert client.image_generation("p", reference_images=["https://img/x.jpg"]) is None
        err = client.get_last_image_error()
        assert err["error_code"] == "portrait_intercept"

    def test_500_retries_then_fails(self):
        client = _make_client(max_retries=1)
        bad_resp = MagicMock(status_code=500)
        bad_resp.text = "internal error"
        bad_resp.json.return_value = {"error": {"message": "internal"}}
        bad_resp.raise_for_status.side_effect = httpx.HTTPStatusError("500", request=MagicMock(), response=bad_resp)
        with (
            patch("packages.shared.ai_client.httpx.post", return_value=bad_resp) as mp,
            patch("packages.shared.ai_client.time.sleep", return_value=None),
        ):
            assert client.image_generation("p") is None
            assert mp.call_count == 2
        err = client.get_last_image_error()
        assert err["error_code"] == "network_error"


# ── 信任链集成单测 ────────────────────────────────────────────────


class TestTrustChainIntegration:
    def test_with_reference_image_triggers_seedream_then_seedance_with_reference_image_role(self, tmp_path):
        client = _make_client()
        captured_calls = []

        seedream_ok = MagicMock(status_code=200)
        seedream_ok.json.return_value = {"data": [{"url": "https://ai.example.com/trusted.png"}]}
        seedream_ok.raise_for_status = MagicMock()
        seedream_ok.text = ""

        task_ok = MagicMock(status_code=200)
        task_ok.json.return_value = {"id": "t-trust"}
        task_ok.raise_for_status = MagicMock()
        task_ok.text = ""

        poll_ok = MagicMock(status_code=200)
        poll_ok.json.return_value = {"status": "succeeded", "content": {"video_url": "https://cdn.example.com/v.mp4"}}
        poll_ok.raise_for_status = MagicMock()

        class FakeStream:
            def __init__(self):
                self._c = [b"OK"]
                self._it = iter(self._c)

            def __enter__(self):
                return self

            def __exit__(self, *a):
                return False

            def raise_for_status(self):
                return None

            def iter_bytes(self, chunk_size=None):
                return self._it

        def fake_post(url, **kwargs):
            captured_calls.append({"url": url, "json": kwargs.get("json")})
            if "/images/generations" in url:
                return seedream_ok
            return task_ok

        fake_uuid = MagicMock()
        fake_uuid.hex = "00000001"
        with (
            patch("packages.shared.ai_client.httpx.post", side_effect=fake_post),
            patch("packages.shared.ai_client.httpx.get", return_value=poll_ok),
            patch("packages.shared.ai_client.httpx.stream", return_value=FakeStream()),
            patch("packages.shared.ai_client.time.sleep", return_value=None),
            patch("packages.shared.ai_client.time.time", side_effect=iter([1000.0, 1001.0, 1002.0, 9e9])),
            patch("packages.shared.ai_client.uuid.uuid4", return_value=fake_uuid),
            patch("packages.shared.ai_client.get_shared_settings") as ms,
        ):
            ms.return_value = MagicMock(
                doubao_video_poll_interval=0,
                doubao_video_timeout=60,
                doubao_video_model="doubao-seedance-2-5-260628",
            )
            out = client.video_generation(
                "人物在海边散步",
                image_url="https://img/raw.jpg",
                duration=5,
                ratio="9:16",
                resolution="720p",
                output_dir=str(tmp_path),
            )
        assert out is not None
        assert len(captured_calls) == 2
        assert "/images/generations" in captured_calls[0]["url"]
        assert captured_calls[0]["json"]["image"] == "https://img/raw.jpg"
        seedance_payload = captured_calls[1]["json"]
        content = seedance_payload["content"]
        img_items = [c for c in content if c.get("type") == "image_url"]
        assert len(img_items) == 1
        assert img_items[0]["image_url"]["url"] == "https://ai.example.com/trusted.png"
        assert img_items[0]["role"] == "reference_image"
        assert seedance_payload["ratio"] == "9:16"

    def test_no_reference_image_skips_seedream(self, tmp_path):
        client = _make_client()
        captured_calls = []

        task_ok = MagicMock(status_code=200)
        task_ok.json.return_value = {"id": "t-t2v"}
        task_ok.raise_for_status = MagicMock()
        poll_ok = MagicMock(status_code=200)
        poll_ok.json.return_value = {"status": "succeeded", "content": {"video_url": "https://cdn.example.com/v.mp4"}}
        poll_ok.raise_for_status = MagicMock()

        class FakeStream:
            def __init__(self):
                self._it = iter([b"OK"])

            def __enter__(self):
                return self

            def __exit__(self, *a):
                return False

            def raise_for_status(self):
                return None

            def iter_bytes(self, chunk_size=None):
                return self._it

        def fake_post(url, **kwargs):
            captured_calls.append({"url": url, "json": kwargs.get("json")})
            return task_ok

        fake_uuid = MagicMock()
        fake_uuid.hex = "00000002"
        with (
            patch("packages.shared.ai_client.httpx.post", side_effect=fake_post),
            patch("packages.shared.ai_client.httpx.get", return_value=poll_ok),
            patch("packages.shared.ai_client.httpx.stream", return_value=FakeStream()),
            patch("packages.shared.ai_client.time.sleep", return_value=None),
            patch("packages.shared.ai_client.time.time", side_effect=iter([1000.0, 1001.0, 9e9])),
            patch("packages.shared.ai_client.uuid.uuid4", return_value=fake_uuid),
            patch("packages.shared.ai_client.get_shared_settings") as ms,
        ):
            ms.return_value = MagicMock(
                doubao_video_poll_interval=0,
                doubao_video_timeout=60,
                doubao_video_model="doubao-seedance-2-5-260628",
            )
            out = client.video_generation("海边日落", duration=5, ratio="9:16", output_dir=str(tmp_path))
        assert out is not None
        assert len(captured_calls) == 1
        assert "/contents/generations/tasks" in captured_calls[0]["url"]
        content = captured_calls[0]["json"]["content"]
        assert all(c.get("type") != "image_url" for c in content)

    def test_seedream_failure_falls_back_to_original_image(self, tmp_path):
        client = _make_client(max_retries=0)
        captured_calls = []

        seedream_fail = MagicMock(status_code=400)
        seedream_fail.text = '{"error":{"code":"QuotaExceeded","message":"quota"}}'
        seedream_fail.json.return_value = {"error": {"code": "QuotaExceeded"}}
        seedream_fail.raise_for_status.side_effect = httpx.HTTPStatusError(
            "q", request=MagicMock(), response=seedream_fail
        )

        task_ok = MagicMock(status_code=200)
        task_ok.json.return_value = {"id": "t-fb"}
        task_ok.raise_for_status = MagicMock()
        poll_ok = MagicMock(status_code=200)
        poll_ok.json.return_value = {"status": "succeeded", "content": {"video_url": "https://cdn.example.com/v.mp4"}}
        poll_ok.raise_for_status = MagicMock()

        class FakeStream:
            def __init__(self):
                self._it = iter([b"OK"])

            def __enter__(self):
                return self

            def __exit__(self, *a):
                return False

            def raise_for_status(self):
                return None

            def iter_bytes(self, chunk_size=None):
                return self._it

        def fake_post(url, **kwargs):
            captured_calls.append({"url": url, "json": kwargs.get("json")})
            if "/images/generations" in url:
                return seedream_fail
            return task_ok

        fake_uuid = MagicMock()
        fake_uuid.hex = "00000003"
        with (
            patch("packages.shared.ai_client.httpx.post", side_effect=fake_post),
            patch("packages.shared.ai_client.httpx.get", return_value=poll_ok),
            patch("packages.shared.ai_client.httpx.stream", return_value=FakeStream()),
            patch("packages.shared.ai_client.time.sleep", return_value=None),
            patch("packages.shared.ai_client.time.time", side_effect=iter([1000.0, 1001.0, 9e9])),
            patch("packages.shared.ai_client.uuid.uuid4", return_value=fake_uuid),
            patch("packages.shared.ai_client.get_shared_settings") as ms,
        ):
            ms.return_value = MagicMock(
                doubao_video_poll_interval=0,
                doubao_video_timeout=60,
                doubao_video_model="doubao-seedance-2-5-260628",
            )
            out = client.video_generation(
                "海边散步",
                image_url="https://img/raw.jpg",
                duration=5,
                ratio="9:16",
                output_dir=str(tmp_path),
            )
        assert out is not None
        assert len(captured_calls) == 2
        seedance_payload = captured_calls[1]["json"]
        content = seedance_payload["content"]
        img_items = [c for c in content if c.get("type") == "image_url"]
        assert len(img_items) == 1
        assert img_items[0]["image_url"]["url"] == "https://img/raw.jpg"
        assert img_items[0]["role"] == "first_frame"
        assert seedance_payload["ratio"] == "adaptive"