From ab99737e03c609ff7b9eb5f2dc62d926d103b785 Mon Sep 17 00:00:00 2001 From: chengma Date: Sat, 11 Jul 2026 12:42:27 +0800 Subject: [PATCH] feat(ai-studio): add hosted image generation jobs --- app/image_studio_generation.py | 392 ++++++++++++++++++++++++++ docs/tasks/T-590.md | 11 +- tests/test_image_studio_generation.py | 282 ++++++++++++++++++ 3 files changed, 683 insertions(+), 2 deletions(-) create mode 100644 app/image_studio_generation.py create mode 100644 tests/test_image_studio_generation.py diff --git a/app/image_studio_generation.py b/app/image_studio_generation.py new file mode 100644 index 0000000..a294005 --- /dev/null +++ b/app/image_studio_generation.py @@ -0,0 +1,392 @@ +"""cmhub hosted generation orchestration for AI image studio jobs.""" + +from __future__ import annotations + +import os +import threading +import time +import urllib.parse +from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait + +from . import ai, appconfig, image_studio +from .version import APP_VERSION + + +MAX_CMHUB_IMAGE_STUDIO_WORKERS = 5 + + +class ImageStudioGenerationError(RuntimeError): + """Raised when AI studio hosted generation cannot complete.""" + + +def _notify(callback, event): + if callback is None: + return + try: + callback(dict(event)) + except Exception: + pass + + +def _runtime(config, cmhub_config_path): + return ai._cmhub_runtime(config, "image", cmhub_config_path) + + +def _generated_kind(job_type): + return "generated_detail" if str(job_type) == "detail" else "generated_main" + + +def _output_path(project, job, image_root): + dirs = image_studio.project_image_dirs(image_root, project) + filename = "job_%06d_%s.jpg" % (int(job.id), _generated_kind(job.job_type)) + return os.path.join(dirs["generated"], filename) + + +def _source_path(source_asset): + path = str(getattr(source_asset, "local_path", "") or "").strip() + if not path or not os.path.isfile(path): + raise ImageStudioGenerationError("源图尚未下载到本地，不能生成") + return path + + +def create_generation_jobs( + project_id, + source_asset_id, + prompt, + count, + *, + job_type="main", + path=None, +): + total = max(0, int(count or 0)) + if total <= 0: + raise ImageStudioGenerationError("生成数量必须大于0") + jobs = [] + for _ in range(total): + jobs.append( + image_studio.create_job( + project_id, + source_asset_id=source_asset_id, + job_type=job_type, + prompt=prompt, + generation_source="cmhub", + provider="cmhub", + path=path, + ) + ) + return jobs + + +def generate_image_jobs( + project_id, + source_asset_id, + prompt, + count, + *, + job_type="main", + aspect_ratio="1:1", + config=None, + cmhub_config_path=None, + path=None, + should_stop=None, + on_event=None, +): + jobs = create_generation_jobs( + project_id, + source_asset_id, + prompt, + count, + job_type=job_type, + path=path, + ) + return run_jobs( + jobs, + aspect_ratio=aspect_ratio, + config=config, + cmhub_config_path=cmhub_config_path, + path=path, + should_stop=should_stop, + on_event=on_event, + ) + + +def resume_image_jobs( + *, + project_id=None, + aspect_ratio="1:1", + config=None, + cmhub_config_path=None, + path=None, + should_stop=None, + on_event=None, +): + jobs = image_studio.list_resumable_jobs(path=path, project_id=project_id) + return run_jobs( + jobs, + aspect_ratio=aspect_ratio, + config=config, + cmhub_config_path=cmhub_config_path, + path=path, + should_stop=should_stop, + on_event=on_event, + ) + + +def run_jobs( + jobs, + *, + aspect_ratio="1:1", + config=None, + cmhub_config_path=None, + path=None, + should_stop=None, + on_event=None, +): + cfg = appconfig.load_config() if config is None else config + job_list = list(jobs or []) + if not job_list: + return {"total": 0, "success": 0, "failed": 0, "cancelled": 0, "jobs": []} + runtime = _runtime(cfg, cmhub_config_path) + ai_cfg = appconfig.ai_config(cfg) + image_root = appconfig.image_dir(cfg) + max_workers = min(MAX_CMHUB_IMAGE_STUDIO_WORKERS, max(1, int(ai_cfg.get("image_concurrency", 1) or 1)), len(job_list)) + should_stop = should_stop or (lambda: False) + lock = threading.Lock() + summary = {"total": len(job_list), "success": 0, "failed": 0, "cancelled": 0, "jobs": []} + + def record(result): + with lock: + summary["jobs"].append(result) + if result.get("status") == "succeeded": + summary["success"] += 1 + elif result.get("status") == "cancelled": + summary["cancelled"] += 1 + else: + summary["failed"] += 1 + + with ThreadPoolExecutor(max_workers=max_workers) as executor: + futures = { + executor.submit( + _run_one_job, + job.id, + runtime, + cfg, + image_root, + aspect_ratio, + path, + should_stop, + on_event, + ): job + for job in job_list + } + while futures: + done, _ = wait(set(futures), return_when=FIRST_COMPLETED) + for future in done: + futures.pop(future) + try: + record(future.result()) + except Exception as exc: + record({"status": "failed", "error": str(exc)}) + if should_stop(): + for future, job in list(futures.items()): + if future.cancel(): + image_studio.update_job_status(job.id, "cancelled", error="用户停止", path=path) + futures.pop(future, None) + record({"job": job, "status": "cancelled", "error": "用户停止"}) + return summary + + +def _run_one_job(job_id, runtime, config, image_root, aspect_ratio, db_path, should_stop, on_event): + job = image_studio.get_job(job_id, path=db_path) + if job is None: + raise ImageStudioGenerationError("AI工场生图任务不存在") + if should_stop(): + updated = image_studio.update_job_status(job.id, "cancelled", error="用户停止", path=db_path) + return {"job": updated, "status": "cancelled", "error": "用户停止"} + project = image_studio.get_project(job.project_id, path=db_path) + source_asset = image_studio.get_asset(job.source_asset_id, path=db_path) + if project is None or source_asset is None: + updated = image_studio.update_job_status(job.id, "failed", error="项目或源图不存在", path=db_path) + return {"job": updated, "status": "failed", "error": "项目或源图不存在"} + try: + image_studio.update_job_status(job.id, "running", path=db_path) + request_result = _submit_or_resume_job( + job, + source_asset, + runtime, + config, + db_path, + should_stop, + on_event, + ) + image_studio.update_job_status(job.id, "running", path=db_path) + request_result = _poll_job(job.id, request_result["task_id"], runtime, request_result, db_path, should_stop, on_event) + out_path = _output_path(project, job, image_root) + saved_path = _download_and_save_job_image(request_result, out_path, config, on_event, job.id) + asset = image_studio.add_asset( + project.id, + _generated_kind(job.job_type), + remote_url=request_result["image_url"], + local_path=saved_path, + aspect_ratio=aspect_ratio, + parent_asset_id=source_asset.id, + prompt=job.prompt, + path=db_path, + ) + updated = image_studio.update_job_status( + job.id, + "succeeded", + output_asset_id=asset.id, + points_balance=request_result.get("points_balance"), + path=db_path, + ) + _notify(on_event, {"job_id": job.id, "step": "job_done", "result": "success"}) + return {"job": updated, "asset": asset, "status": "succeeded"} + except Exception as exc: + status = "cancelled" if "停止" in str(exc) else "failed" + updated = image_studio.update_job_status(job.id, status, error=str(exc), path=db_path) + _notify(on_event, {"job_id": job.id, "step": "job_done", "result": status, "detail": str(exc)}) + return {"job": updated, "status": status, "error": str(exc)} + + +def _submit_or_resume_job(job, source_asset, runtime, config, db_path, should_stop, on_event): + if job.task_id: + _notify(on_event, {"job_id": job.id, "step": "cover_request", "result": "resume", "task_id": job.task_id}) + return _request_result(job.task_id, runtime, config) + _raise_if_stopped(should_stop) + source_path = _source_path(source_asset) + ai_cfg = appconfig.ai_config(config) + resolution = str(ai_cfg.get("resolution", "1k") or "1k") + payload = { + "prompt": str(job.prompt or ""), + "model": runtime["alias"], + "image_base64": ai._image_data_url(source_path), + "resolution": ai._normalize_cmhub_resolution(resolution), + "aspect_ratio": "1:1", + } + _notify(on_event, {"job_id": job.id, "step": "cover_submit", "result": "start"}) + data = ai._cmhub_call_with_retry( + "POST", + appconfig.cmhub_request_url(runtime["base_url"], "/api/v1/generate/image/tasks"), + runtime["api_key"], + payload=payload, + connect_timeout=runtime["connect_timeout"], + read_timeout=ai.CMHUB_IMAGE_SUBMIT_READ_TIMEOUT_SECONDS, + attempts=max(1, int(ai_cfg.get("retry", 0) or 0) + 1), + headers_extra={ + "Idempotency-Key": job.task_key, + "X-Client-Version": str(APP_VERSION), + }, + ) + image_task_id = str(data.get("task_id") or "").strip() + if not image_task_id: + raise ImageStudioGenerationError("cmhub 生图任务提交返回格式错误") + submitted = image_studio.set_job_submitted( + job.id, + image_task_id, + call_id=data.get("call_id"), + points_cost=data.get("points_cost"), + points_balance=data.get("points_balance"), + path=db_path, + ) + _notify( + on_event, + { + "job_id": job.id, + "step": "cover_submit", + "result": "success", + "task_id": submitted.task_id, + "call_id": submitted.call_id, + "points_cost": submitted.points_cost, + "points_balance": submitted.points_balance, + }, + ) + return _request_result(image_task_id, runtime, config, points_balance=data.get("points_balance")) + + +def _request_result(task_id, runtime, config, points_balance=None): + ai_cfg = appconfig.ai_config(config) + resolution = str(ai_cfg.get("resolution", "1k") or "1k") + return { + "task_id": task_id, + "connect_timeout": runtime["connect_timeout"], + "read_timeout": ai.CMHUB_IMAGE_READ_TIMEOUT_SECONDS, + "resolution": resolution, + "quality": ai._jpg_quality(ai_cfg.get("jpg_quality", 90)), + "use_system_proxy": runtime["use_system_proxy"], + "download_with_curl": runtime["download_with_curl"], + "points_balance": points_balance, + } + + +def _poll_job(job_id, task_id, runtime, request_result, db_path, should_stop, on_event): + poll_url = appconfig.cmhub_request_url( + runtime["base_url"], + "/api/v1/generate/image/tasks/%s" % urllib.parse.quote(str(task_id), safe=""), + ) + started = time.perf_counter() + deadline = started + max(1, int(ai.CMHUB_IMAGE_READ_TIMEOUT_SECONDS)) + poll_index = 0 + while True: + _raise_if_stopped(should_stop) + if time.perf_counter() >= deadline: + raise ImageStudioGenerationError("等待 cmhub 生图任务完成超时，下次可继续查询") + _notify(on_event, {"job_id": job_id, "step": "cover_poll", "result": "start", "task_id": task_id}) + data = ai._cmhub_call_once( + "GET", + poll_url, + runtime["api_key"], + payload=None, + connect_timeout=runtime["connect_timeout"], + read_timeout=ai.CMHUB_IMAGE_POLL_READ_TIMEOUT_SECONDS, + headers_extra={"X-Client-Version": str(APP_VERSION)}, + ) + status = str(data.get("status") or "").strip().lower() + if status in {"queued", "running"}: + _notify(on_event, {"job_id": job_id, "step": "cover_poll", "result": status, "task_id": task_id}) + ai._sleep_cmhub_poll(poll_index, should_stop) + poll_index += 1 + continue + if status == "succeeded": + image_url = ai._extract_cmhub_image_url(data, runtime["base_url"]) + if not image_url: + raise ImageStudioGenerationError("cmhub 生图任务成功但没有图片地址") + merged = dict(request_result) + merged["image_url"] = image_url + merged["points_balance"] = data.get("points_balance", merged.get("points_balance")) + _notify(on_event, {"job_id": job_id, "step": "cover_poll", "result": "success", "task_id": task_id}) + return merged + if status in {"failed", "expired"}: + error = data.get("error") if isinstance(data.get("error"), dict) else {} + message = str(error.get("message") or error.get("code") or status) + image_studio.update_job_status(job_id, status, error=message, path=db_path) + raise ImageStudioGenerationError(message) + raise ImageStudioGenerationError("cmhub 生图任务状态返回格式错误") + + +def _download_and_save_job_image(request_result, out_path, config, on_event, job_id): + _notify(on_event, {"job_id": job_id, "step": "cover_download", "result": "start"}) + image_bytes, _ = ai._download_cmhub_image_with_retry( + request_result["image_url"], + connect_timeout=request_result["connect_timeout"], + read_timeout=request_result["read_timeout"], + use_system_proxy=request_result.get("use_system_proxy", False), + download_with_curl=request_result.get("download_with_curl", "auto"), + ) + saved_path = ai._save_jpeg( + image_bytes, + out_path, + request_result.get("resolution") or appconfig.ai_config(config).get("resolution", "1k"), + request_result.get("quality") or appconfig.ai_config(config).get("jpg_quality", 90), + ) + _notify(on_event, {"job_id": job_id, "step": "cover_download", "result": "success"}) + return saved_path + + +def _raise_if_stopped(should_stop): + try: + stopped = bool(should_stop and should_stop()) + except Exception: + stopped = False + if stopped: + raise ImageStudioGenerationError("用户已停止，已提交任务可稍后继续查询") diff --git a/docs/tasks/T-590.md b/docs/tasks/T-590.md index a317710..d488e0e 100644 --- a/docs/tasks/T-590.md +++ b/docs/tasks/T-590.md @@ -3,7 +3,7 @@ id: T-590 title: AI工场托管模式多图异步生成编排与重启续查 phase: 7 deps: [T-586, T-564] -status: TODO +status: DONE created: 2026-07-11 --- @@ -37,4 +37,11 @@ AI工场一次可生成 N 张，且每张都要独立计费、失败、重试、 ## 执行记录 -（完成后记录抽取边界、任务生命周期和回归测试。） +- 2026-07-11：完成 AI工场 cmhub 托管多图异步生成 service。 + - 新增 `app/image_studio_generation.py`，不改现有② AI生成路径；AI工场单独使用 `image_studio_jobs/assets` 持久化 job/key/task_id/计费/状态和输出 asset。 + - `generate_image_jobs()` 对同一源 asset + 完整提示词创建 N 条独立 job；每条 job 单独提交 `/api/v1/generate/image/tasks`，使用自身 `task_key` 作为 `Idempotency-Key`，submit 成功后立即写 `task_id/call_id/points_cost/points_balance`。 + - `resume_image_jobs()` 对已有 `task_id` 的 `submitted/running` job 只 GET poll，不再 POST，避免重启续查重复扣点；poll 成功后下载保存 JPEG 到项目 `generated/` 并创建 `generated_main/generated_detail` asset，关联父源图、比例、提示词和远程 `image_url`。 + - 单张失败只更新该 job 为 `failed`，不回滚其他成功图；下载失败发生在 poll 成功之后，不重新 submit、不创建输出 asset；停止语义保留已提交 job 的 task_id，未开始/可取消 job 标记 `cancelled`。 + - 新增 `tests/test_image_studio_generation.py` 覆盖 N=3 独立 job/key/task、submit 后 poll 前已落库、续查不 POST、下载失败不重提交/不创建 asset、服务端单张 failed 不影响其他成功。 + - 验证：在只套用 T-590 diff 的 clean worktree 中运行 `py -3.10 -m unittest tests.test_image_studio_generation tests.test_ai tests.test_db`（66 tests）、`python -m ruff check app tests main.py`、`py -3.10 -m compileall app main.py`、`py -3.10 -m unittest discover -s tests`（360 tests）和 `git diff --check`，全部通过。 + - 本任务未接 GUI，也未打真实 cmhub 收费接口；T-591 接 UI 时应在 worker/QThread 中调用该 service，避免阻塞主线程。 diff --git a/tests/test_image_studio_generation.py b/tests/test_image_studio_generation.py new file mode 100644 index 0000000..b55bfed --- /dev/null +++ b/tests/test_image_studio_generation.py @@ -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()