cmshoppe/tests/test_image_studio_generation.py

import io
import os
import sys
import threading
import time
import unittest
from concurrent.futures import CancelledError
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.assertEqual(1, len(submitted_payloads[0]["images"]))
            self.assertIn("image_base64", submitted_payloads[0]["images"][0])
            self.assertNotIn("image_base64", submitted_payloads[0])

        self.assert_removed(temp_dir)

    def test_build_cmhub_images_keeps_order_and_limits_to_eight_inputs(self):
        with self.make_temp_dir() as temp_dir:
            cfg, project, source = self._project_source(temp_dir)
            references = []
            for index in range(2, 10):
                path = os.path.join(temp_dir, "source-%d.png" % index)
                with open(path, "wb") as fh:
                    fh.write(self._png_bytes())
                references.append(
                    image_studio.add_asset(
                        project.id,
                        image_studio.ASSET_KIND_ORIGINAL,
                        local_path=path,
                        source_order=index,
                        path=cfg["db_path"],
                    )
                )

            with mock.patch(
                "app.image_studio_generation.ai._image_data_url",
                side_effect=lambda path: "data:image/png;base64,%s" % os.path.basename(path),
            ):
                images, omitted_count = image_studio_generation._build_cmhub_images(
                    source,
                    references,
                )

            self.assertEqual(8, len(images))
            self.assertEqual(1, omitted_count)
            self.assertTrue(images[0]["image_base64"].endswith("source.png"))
            self.assertTrue(images[-1]["image_base64"].endswith("source-8.png"))

        self.assert_removed(temp_dir)

    def test_build_cmhub_images_rejects_oversized_total_payload(self):
        with self.make_temp_dir() as temp_dir:
            _, _, source = self._project_source(temp_dir)
            with mock.patch(
                "app.image_studio_generation.ai._image_data_url",
                return_value="x" * 32,
            ), mock.patch.object(
                image_studio_generation,
                "CMHUB_IMAGE_STUDIO_MAX_TOTAL_INPUT_BYTES",
                16,
            ):
                with self.assertRaisesRegex(
                    image_studio_generation.ImageStudioGenerationError,
                    "总大小超过32MiB",
                ):
                    image_studio_generation._build_cmhub_images(source)

        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,
                should_stop=None,
            ):
                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_download_cancel_maps_to_cancelled_and_keeps_resume_action(self):
        with self.make_temp_dir() as temp_dir:
            cfg, project, source = self._project_source(temp_dir)
            stopped = {"value": False}

            def fake_download(*args, **kwargs):
                self.assertIsNotNone(kwargs.get("should_stop"))
                stopped["value"] = True
                raise CancelledError()

            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-cancel-download", "status": "queued"},
            ), mock.patch(
                "app.image_studio_generation.ai._cmhub_call_once",
                return_value={
                    "task_id": "cmhub-cancel-download",
                    "status": "succeeded",
                    "result": {
                        "image_url": "https://cdn.example.com/cancel-download.png"
                    },
                },
            ), mock.patch(
                "app.image_studio_generation.ai._download_cmhub_image_with_retry",
                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"])
            job = summary["jobs"][0]["job"]
            stored = image_studio.get_job(job.id, path=cfg["db_path"])
            self.assertEqual("cancelled", stored.status)
            self.assertEqual(image_studio.JOB_RECOVERY_RESUME, stored.recovery_action)
            self.assertEqual(
                [],
                image_studio.list_assets(
                    project.id,
                    kind="generated_main",
                    path=cfg["db_path"],
                ),
            )

        self.assert_removed(temp_dir)

    def test_run_jobs_cancels_queued_future_before_running_slot_is_released(self):
        with self.make_temp_dir() as temp_dir:
            cfg, project, source = self._project_source(temp_dir)
            cfg["ai"]["image_concurrency"] = 1
            jobs = [
                image_studio.create_job(
                    project.id,
                    source_asset_id=source.id,
                    job_type="main",
                    prompt="排队停止测试",
                    path=cfg["db_path"],
                )
                for _ in range(2)
            ]
            started_ids = []
            first_started = threading.Event()
            release_first = threading.Event()
            stop_requested = threading.Event()
            result_holder = {}

            def fake_run_one(*args):
                job_id = int(args[0])
                started_ids.append(job_id)
                if len(started_ids) == 1:
                    first_started.set()
                    release_first.wait(timeout=3)
                updated = image_studio.update_job_status(
                    job_id,
                    "cancelled",
                    error="用户停止",
                    recovery_action=image_studio.JOB_RECOVERY_REGENERATE,
                    path=cfg["db_path"],
                )
                return {
                    "job": updated,
                    "status": "cancelled",
                    "error": "用户停止",
                }

            def run():
                result_holder["summary"] = image_studio_generation.run_jobs(
                    jobs,
                    config=cfg,
                    path=cfg["db_path"],
                    should_stop=stop_requested.is_set,
                )

            with mock.patch(
                "app.image_studio_generation._runtime",
                return_value=self._runtime(),
            ), mock.patch(
                "app.image_studio_generation._run_one_job_with_global_slot",
                side_effect=fake_run_one,
            ):
                thread = threading.Thread(target=run)
                thread.start()
                self.assertTrue(first_started.wait(timeout=2))
                stop_requested.set()
                time.sleep(0.35)
                release_first.set()
                thread.join(timeout=3)

            self.assertFalse(thread.is_alive())
            self.assertEqual([jobs[0].id], started_ids)
            summary = result_holder["summary"]
            self.assertEqual(2, summary["cancelled"])
            self.assertEqual(2, len(summary["jobs"]))
            self.assertEqual(
                ["cancelled", "cancelled"],
                [
                    image_studio.get_job(job.id, path=cfg["db_path"]).status
                    for job in jobs
                ],
            )

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