life-echo/api/tests/test_memoir_image_provider.py

import unittest
from unittest.mock import Mock, patch

from api.services.memoir_images.provider import LiblibImageProvider


def _make_provider(http_client=None):
    return LiblibImageProvider(
        http_client=http_client or Mock(),
        access_key="test-ak",
        secret_key="test-sk",
        base_url="https://openapi.liblibai.cloud",
        template_uuid="tpl-uuid",
    )


class LiblibSignatureTest(unittest.TestCase):
    def test_sign_returns_url_with_auth_params(self):
        provider = _make_provider()
        url = provider._sign("/api/generate/webui/text2img/ultra")
        self.assertIn("AccessKey=test-ak", url)
        self.assertIn("Signature=", url)
        self.assertIn("Timestamp=", url)
        self.assertIn("SignatureNonce=", url)
        self.assertTrue(url.startswith("https://openapi.liblibai.cloud/api/generate/webui/text2img/ultra?"))


class SubmitGenerationTest(unittest.TestCase):
    def test_submit_returns_processing_with_job_id(self):
        http_client = Mock()
        resp = Mock()
        resp.json.return_value = {
            "code": 0,
            "data": {"generateUuid": "uuid-abc"},
        }
        resp.raise_for_status = Mock()
        http_client.post.return_value = resp

        provider = _make_provider(http_client)
        job = provider.submit_generation(prompt="a cat", size="1024x1024", style="watercolor")

        self.assertEqual(job["status"], "processing")
        self.assertEqual(job["job_id"], "uuid-abc")
        self.assertIsNone(job["image_url"])

        call_kwargs = http_client.post.call_args
        body = call_kwargs.kwargs.get("json") or call_kwargs[1].get("json")
        self.assertEqual(body["templateUuid"], "tpl-uuid")
        self.assertEqual(body["generateParams"]["prompt"], "a cat")
        self.assertEqual(body["generateParams"]["aspectRatio"], "square")

    def test_submit_raises_on_error_code(self):
        http_client = Mock()
        resp = Mock()
        resp.json.return_value = {"code": 100000, "msg": "param error"}
        resp.raise_for_status = Mock()
        http_client.post.return_value = resp

        provider = _make_provider(http_client)
        with self.assertRaises(RuntimeError):
            provider.submit_generation(prompt="a cat", size="1024x1024", style="watercolor")


class PollUntilCompleteTest(unittest.TestCase):
    def test_returns_completed_on_status_5(self):
        http_client = Mock()
        pending_resp = Mock()
        pending_resp.json.return_value = {
            "code": 0,
            "data": {"generateStatus": 2, "images": []},
        }
        pending_resp.raise_for_status = Mock()

        success_resp = Mock()
        success_resp.json.return_value = {
            "code": 0,
            "data": {
                "generateStatus": 5,
                "images": [{"imageUrl": "https://cdn.example.com/1.png", "auditStatus": 3}],
            },
        }
        success_resp.raise_for_status = Mock()
        http_client.post.side_effect = [pending_resp, success_resp]

        provider = _make_provider(http_client)
        job = provider.poll_until_complete(
            {"status": "processing", "job_id": "uuid-abc"},
            poll_interval_seconds=0,
            max_attempts=3,
        )

        self.assertEqual(job["status"], "completed")
        self.assertEqual(job["image_url"], "https://cdn.example.com/1.png")
        self.assertEqual(job["job_id"], "uuid-abc")

    def test_raises_on_status_6_failure(self):
        http_client = Mock()
        resp = Mock()
        resp.json.return_value = {
            "code": 0,
            "data": {"generateStatus": 6, "generateMsg": "content violation"},
        }
        resp.raise_for_status = Mock()
        http_client.post.return_value = resp

        provider = _make_provider(http_client)
        with self.assertRaises(RuntimeError, msg="content violation"):
            provider.poll_until_complete(
                {"status": "processing", "job_id": "uuid-abc"},
                poll_interval_seconds=0,
                max_attempts=2,
            )

    def test_raises_timeout_after_max_attempts(self):
        http_client = Mock()
        resp = Mock()
        resp.json.return_value = {
            "code": 0,
            "data": {"generateStatus": 2, "images": []},
        }
        resp.raise_for_status = Mock()
        http_client.post.return_value = resp

        provider = _make_provider(http_client)
        with self.assertRaises(TimeoutError):
            provider.poll_until_complete(
                {"status": "processing", "job_id": "uuid-abc"},
                poll_interval_seconds=0,
                max_attempts=2,
            )


class DownloadImageTest(unittest.TestCase):
    def test_download_fetches_binary_payload(self):
        http_client = Mock()
        resp = Mock()
        resp.content = b"png-bytes"
        resp.raise_for_status = Mock()
        http_client.get.return_value = resp

        provider = _make_provider(http_client)
        payload = provider.download_image({"image_url": "https://cdn.example.com/1.png"})

        self.assertEqual(payload, b"png-bytes")