fix(product-suite): finalize and cancel generation

tests/test_image_studio_generation.py

@@ -1,7 +1,10 @@
 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__))
@@ -193,7 +196,14 @@ class ImageStudioGenerationTests(TempDirMixin, unittest.TestCase):
             stopped = {"value": False}
             saved_paths = []
 
-            def fake_download(request_result, out_path, config, on_event, job_id):
+            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())
@@ -235,6 +245,138 @@ class ImageStudioGenerationTests(TempDirMixin, unittest.TestCase):
 
         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)