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")