T-565 handle relative cmhub image urls

tests/test_ai.py

@@ -1155,6 +1155,7 @@ class AITests(TempDirMixin, unittest.TestCase):
             ]
             poll_db_values = []
             headers_seen = []
+            downloads = []
 
             def fake_request(method, url, **kwargs):
                 if str(method).upper() == "POST":
@@ -1174,11 +1175,12 @@ class AITests(TempDirMixin, unittest.TestCase):
                     {
                         "task_id": "cmhub-task-1",
                         "status": "succeeded",
-                        "result": {"image_url": "https://cdn.example.com/generated.png"},
+                        "result": {"image_url": "/generated/images/generated.png"},
                     }
                 )
 
             def fake_get(url, **kwargs):
+                downloads.append((url, kwargs))
                 return _RequestsResponse(content=generated_png)
 
             with mock.patch.object(ai._cmhub_session(), "request", side_effect=fake_request), \
@@ -1198,6 +1200,10 @@ class AITests(TempDirMixin, unittest.TestCase):
             self.assertEqual(["cmhub-task-1"], poll_db_values)
             self.assertTrue(headers_seen[0]["Idempotency-Key"].startswith("cmshopee-task-"))
             self.assertTrue(headers_seen[0]["X-Client-Version"])
+            self.assertEqual(
+                "https://cmhub.example.com/generated/images/generated.png",
+                downloads[0][0],
+            )
             updated = db.get_task(task.id, path=cfg["db_path"])
             self.assertEqual("cmhub-task-1", updated.image_task_id)
             self.assertTrue(updated.image_task_key)
@@ -1205,6 +1211,23 @@ class AITests(TempDirMixin, unittest.TestCase):
 
         self.assert_removed(temp_dir)
 
+    def test_cmhub_async_image_url_extractor_accepts_image_urls_list(self):
+        data = {
+            "task_id": "cmhub-task-1",
+            "status": "succeeded",
+            "result": {
+                "image_urls": [
+                    "https://cdn.example.com/generated-a.png",
+                    "https://cdn.example.com/generated-b.png",
+                ]
+            },
+        }
+
+        self.assertEqual(
+            "https://cdn.example.com/generated-a.png",
+            ai._extract_cmhub_image_url(data, "https://cmhub.example.com"),
+        )
+
     def test_generate_batch_cmhub_resumes_existing_image_task_without_submit(self):
         try:
             from PIL import Image  # noqa: F401