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_generate_image_jobs_sends_selected_aspect_ratio(self): with self.make_temp_dir() as temp_dir: cfg, project, source = self._project_source(temp_dir) submitted_payloads = [] def fake_submit(method, url, api_key, **kwargs): submitted_payloads.append(dict(kwargs["payload"])) return {"task_id": "cmhub-ratio", "status": "queued"} 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", return_value={ "task_id": "cmhub-ratio", "status": "succeeded", "result": {"image_url": "https://cdn.example.com/ratio.png"}, }, ), \ 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, "比例测试", 1, aspect_ratio="3:4", config=cfg, path=cfg["db_path"], ) self.assertEqual(1, summary["success"]) self.assertEqual("3:4", submitted_payloads[0]["aspect_ratio"]) self.assert_removed(temp_dir) def test_stop_after_download_discards_temporary_result(self): with self.make_temp_dir() as temp_dir: cfg, project, source = self._project_source(temp_dir) stopped = {"value": False} saved_paths = [] def fake_download(request_result, out_path, config, on_event, job_id): os.makedirs(os.path.dirname(out_path), exist_ok=True) with open(out_path, "wb") as fh: fh.write(self._png_bytes()) saved_paths.append(out_path) stopped["value"] = True return out_path with mock.patch("app.image_studio_generation._runtime", return_value=self._runtime()), \ mock.patch( "app.image_studio_generation.ai._cmhub_call_with_retry", return_value={"task_id": "cmhub-stop", "status": "queued"}, ), \ mock.patch( "app.image_studio_generation.ai._cmhub_call_once", return_value={ "task_id": "cmhub-stop", "status": "succeeded", "result": {"image_url": "https://cdn.example.com/stop.png"}, }, ), \ mock.patch( "app.image_studio_generation._download_and_save_job_image", side_effect=fake_download, ): summary = image_studio_generation.generate_image_jobs( project.id, source.id, "停止测试", 1, config=cfg, path=cfg["db_path"], should_stop=lambda: stopped["value"], ) self.assertEqual(1, summary["cancelled"]) self.assertEqual([], image_studio.list_assets(project.id, kind="generated_main", path=cfg["db_path"])) self.assertEqual(1, len(saved_paths)) self.assertFalse(os.path.exists(saved_paths[0])) 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_resume_failed_download_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-download", call_id="call-download", points_cost=2, points_balance=88, path=cfg["db_path"], ) image_studio.update_job_status( job.id, "failed", error="下载新封面失败", path=cfg["db_path"], ) failed_job = image_studio.get_job(job.id, path=cfg["db_path"]) self.assertEqual(image_studio.JOB_RECOVERY_RESUME, failed_job.recovery_action) resumable = image_studio.list_resumable_jobs( path=cfg["db_path"], project_id=project.id, include_failed_downloads=True, ) self.assertEqual([job.id], [item.id for item in resumable]) def fake_poll(method, url, api_key, **kwargs): self.assertEqual("GET", method) return { "task_id": "cmhub-task-download", "status": "succeeded", "result": {"image_url": "https://cdn.example.com/recovered.png"}, "points_balance": 88, } 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-download", updated.task_id) self.assertEqual("call-download", updated.call_id) self.assertEqual(image_studio.JOB_RECOVERY_NONE, updated.recovery_action) assets = image_studio.list_assets(project.id, kind="generated_main", path=cfg["db_path"]) self.assertEqual(1, len(assets)) 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)) failed_jobs = [ job for job in image_studio.list_resumable_jobs( path=cfg["db_path"], project_id=project.id, include_failed_downloads=True, ) if job.status == "failed" ] self.assertEqual([], failed_jobs) self.assert_removed(temp_dir) if __name__ == "__main__": unittest.main()