feat(ai-studio): add hosted image generation jobs

2026-07-11 12:42:27 +08:00
parent 3ebefab91e
commit ab99737e03
3 changed files with 683 additions and 2 deletions
@@ -0,0 +1,282 @@
+import io
+import os
+import sys
+import unittest
+from unittest import mock
+
+sys.path.insert(0, os.path.dirname(__file__))
+
+from _helpers import TempDirMixin
+
+from app import db, image_studio, image_studio_generation
+
+
+class ImageStudioGenerationTests(TempDirMixin, unittest.TestCase):
+    def _png_bytes(self):
+        from PIL import Image
+
+        output = io.BytesIO()
+        Image.new("RGB", (32, 32), (120, 80, 160)).save(output, format="PNG")
+        return output.getvalue()
+
+    def _config(self, temp_dir):
+        return {
+            "data_dir": temp_dir,
+            "db_path": os.path.join(temp_dir, "cmshopee.db"),
+            "image_dir": os.path.join(temp_dir, "images"),
+            "ai": {
+                "backend": "cmhub",
+                "image_concurrency": 1,
+                "retry": 0,
+                "resolution": "1k",
+                "jpg_quality": 90,
+                "cmhub": {
+                    "base_url": "https://cmhub.example.com",
+                    "image_alias": "image-hd",
+                    "connect_timeout": 3,
+                    "download_with_curl": "false",
+                },
+            },
+        }
+
+    def _project_source(self, temp_dir):
+        cfg = self._config(temp_dir)
+        db.init_db(cfg["db_path"])
+        project = image_studio.create_or_get_project(
+            account_alias="alias",
+            account_slug="alias_slug",
+            item_id="51100639510",
+            path=cfg["db_path"],
+        )
+        source_path = os.path.join(temp_dir, "source.png")
+        with open(source_path, "wb") as fh:
+            fh.write(self._png_bytes())
+        source = image_studio.add_asset(
+            project.id,
+            image_studio.ASSET_KIND_ORIGINAL,
+            remote_url="https://cdn.example.com/source.png",
+            local_path=source_path,
+            source_order=1,
+            path=cfg["db_path"],
+        )
+        return cfg, project, source
+
+    def _runtime(self):
+        return {
+            "base_url": "https://cmhub.example.com",
+            "api_key": "sk-test",
+            "alias": "image-hd",
+            "connect_timeout": 3,
+            "use_system_proxy": False,
+            "download_with_curl": "false",
+        }
+
+    def test_generate_image_jobs_creates_independent_jobs_and_assets(self):
+        with self.make_temp_dir() as temp_dir:
+            cfg, project, source = self._project_source(temp_dir)
+            submitted = []
+            poll_seen_persisted = []
+
+            def fake_submit(method, url, api_key, **kwargs):
+                self.assertEqual("POST", method)
+                task_id = "cmhub-task-%d" % (len(submitted) + 1)
+                submitted.append((task_id, kwargs["headers_extra"]["Idempotency-Key"]))
+                return {
+                    "task_id": task_id,
+                    "status": "queued",
+                    "call_id": "call-%d" % len(submitted),
+                    "points_cost": 2,
+                    "points_balance": 100 - len(submitted) * 2,
+                }
+
+            def fake_poll(method, url, api_key, **kwargs):
+                self.assertEqual("GET", method)
+                task_id = url.rsplit("/", 1)[-1]
+                jobs = image_studio.list_resumable_jobs(path=cfg["db_path"], project_id=project.id)
+                poll_seen_persisted.append(any(job.task_id == task_id for job in jobs))
+                return {
+                    "task_id": task_id,
+                    "status": "succeeded",
+                    "result": {"image_url": "https://cdn.example.com/%s.png" % task_id},
+                }
+
+            with mock.patch("app.image_studio_generation._runtime", return_value=self._runtime()), \
+                    mock.patch(
+                        "app.image_studio_generation.ai._cmhub_call_with_retry",
+                        side_effect=fake_submit,
+                    ), \
+                    mock.patch(
+                        "app.image_studio_generation.ai._cmhub_call_once",
+                        side_effect=fake_poll,
+                    ), \
+                    mock.patch(
+                        "app.image_studio_generation.ai._download_cmhub_image_with_retry",
+                        return_value=(self._png_bytes(), 0.1),
+                    ), \
+                    mock.patch("app.image_studio_generation.ai._sleep_cmhub_poll"):
+                summary = image_studio_generation.generate_image_jobs(
+                    project.id,
+                    source.id,
+                    "完整提示词",
+                    3,
+                    config=cfg,
+                    path=cfg["db_path"],
+                )
+
+            self.assertEqual(3, summary["success"])
+            self.assertEqual(0, summary["failed"])
+            self.assertEqual(3, len(submitted))
+            self.assertEqual(3, len({key for _, key in submitted}))
+            self.assertEqual([True, True, True], poll_seen_persisted)
+            jobs = image_studio.list_resumable_jobs(path=cfg["db_path"], project_id=project.id)
+            self.assertEqual([], jobs)
+            all_jobs = [
+                image_studio.get_job(result["job"].id, path=cfg["db_path"])
+                for result in summary["jobs"]
+            ]
+            self.assertEqual(["succeeded", "succeeded", "succeeded"], [job.status for job in all_jobs])
+            self.assertTrue(all(job.task_id for job in all_jobs))
+            self.assertTrue(all(job.task_key for job in all_jobs))
+            self.assertEqual([2, 2, 2], [job.points_cost for job in all_jobs])
+            assets = image_studio.list_assets(project.id, kind="generated_main", path=cfg["db_path"])
+            self.assertEqual(3, len(assets))
+            self.assertTrue(all(os.path.isfile(asset.local_path) for asset in assets))
+            self.assertTrue(all(asset.parent_asset_id == source.id for asset in assets))
+            self.assertTrue(all(asset.prompt == "完整提示词" for asset in assets))
+
+        self.assert_removed(temp_dir)
+
+    def test_resume_existing_job_polls_without_new_submit(self):
+        with self.make_temp_dir() as temp_dir:
+            cfg, project, source = self._project_source(temp_dir)
+            job = image_studio.create_job(
+                project.id,
+                source_asset_id=source.id,
+                job_type="main",
+                prompt="续查提示词",
+                path=cfg["db_path"],
+            )
+            image_studio.set_job_submitted(job.id, "cmhub-task-resume", path=cfg["db_path"])
+
+            def fake_poll(method, url, api_key, **kwargs):
+                return {
+                    "task_id": "cmhub-task-resume",
+                    "status": "succeeded",
+                    "result": {"image_url": "https://cdn.example.com/resume.png"},
+                }
+
+            with mock.patch("app.image_studio_generation._runtime", return_value=self._runtime()), \
+                    mock.patch("app.image_studio_generation.ai._cmhub_call_with_retry") as submit, \
+                    mock.patch("app.image_studio_generation.ai._cmhub_call_once", side_effect=fake_poll), \
+                    mock.patch(
+                        "app.image_studio_generation.ai._download_cmhub_image_with_retry",
+                        return_value=(self._png_bytes(), 0.1),
+                    ):
+                summary = image_studio_generation.resume_image_jobs(
+                    project_id=project.id,
+                    config=cfg,
+                    path=cfg["db_path"],
+                )
+
+            self.assertEqual(1, summary["success"])
+            submit.assert_not_called()
+            updated = image_studio.get_job(job.id, path=cfg["db_path"])
+            self.assertEqual("succeeded", updated.status)
+            self.assertEqual("cmhub-task-resume", updated.task_id)
+
+        self.assert_removed(temp_dir)
+
+    def test_download_failure_does_not_submit_again_or_create_asset(self):
+        with self.make_temp_dir() as temp_dir:
+            cfg, project, source = self._project_source(temp_dir)
+            submit_count = 0
+
+            def fake_submit(method, url, api_key, **kwargs):
+                nonlocal submit_count
+                submit_count += 1
+                return {"task_id": "cmhub-task-1", "status": "queued"}
+
+            def fake_poll(method, url, api_key, **kwargs):
+                return {
+                    "task_id": "cmhub-task-1",
+                    "status": "succeeded",
+                    "result": {"image_url": "https://cdn.example.com/fail.png"},
+                }
+
+            with mock.patch("app.image_studio_generation._runtime", return_value=self._runtime()), \
+                    mock.patch(
+                        "app.image_studio_generation.ai._cmhub_call_with_retry",
+                        side_effect=fake_submit,
+                    ), \
+                    mock.patch("app.image_studio_generation.ai._cmhub_call_once", side_effect=fake_poll), \
+                    mock.patch(
+                        "app.image_studio_generation.ai._download_cmhub_image_with_retry",
+                        side_effect=RuntimeError("download failed"),
+                    ):
+                summary = image_studio_generation.generate_image_jobs(
+                    project.id,
+                    source.id,
+                    "完整提示词",
+                    1,
+                    config=cfg,
+                    path=cfg["db_path"],
+                )
+
+            self.assertEqual(0, summary["success"])
+            self.assertEqual(1, summary["failed"])
+            self.assertEqual(1, submit_count)
+            jobs = [result["job"] for result in summary["jobs"]]
+            self.assertEqual("failed", image_studio.get_job(jobs[0].id, path=cfg["db_path"]).status)
+            self.assertEqual([], image_studio.list_assets(project.id, kind="generated_main", path=cfg["db_path"]))
+
+        self.assert_removed(temp_dir)
+
+    def test_failed_cmhub_task_marks_only_that_job_failed(self):
+        with self.make_temp_dir() as temp_dir:
+            cfg, project, source = self._project_source(temp_dir)
+            submitted = []
+
+            def fake_submit(method, url, api_key, **kwargs):
+                task_id = "cmhub-task-%d" % (len(submitted) + 1)
+                submitted.append(task_id)
+                return {"task_id": task_id, "status": "queued"}
+
+            def fake_poll(method, url, api_key, **kwargs):
+                task_id = url.rsplit("/", 1)[-1]
+                if task_id.endswith("-2"):
+                    return {"task_id": task_id, "status": "failed", "error": {"message": "上游失败"}}
+                return {
+                    "task_id": task_id,
+                    "status": "succeeded",
+                    "result": {"image_url": "https://cdn.example.com/%s.png" % task_id},
+                }
+
+            with mock.patch("app.image_studio_generation._runtime", return_value=self._runtime()), \
+                    mock.patch(
+                        "app.image_studio_generation.ai._cmhub_call_with_retry",
+                        side_effect=fake_submit,
+                    ), \
+                    mock.patch("app.image_studio_generation.ai._cmhub_call_once", side_effect=fake_poll), \
+                    mock.patch(
+                        "app.image_studio_generation.ai._download_cmhub_image_with_retry",
+                        return_value=(self._png_bytes(), 0.1),
+                    ):
+                summary = image_studio_generation.generate_image_jobs(
+                    project.id,
+                    source.id,
+                    "完整提示词",
+                    2,
+                    config=cfg,
+                    path=cfg["db_path"],
+                )
+
+            self.assertEqual(1, summary["success"])
+            self.assertEqual(1, summary["failed"])
+            assets = image_studio.list_assets(project.id, kind="generated_main", path=cfg["db_path"])
+            self.assertEqual(1, len(assets))
+
+        self.assert_removed(temp_dir)
+
+
+if __name__ == "__main__":
+    unittest.main()