feat(product-suite): make image pulls cancellable

tests/test_image_studio.py

@@ -1060,6 +1060,95 @@ class ImageStudioTests(TempDirMixin, unittest.TestCase):
 
         self.assert_removed(temp_dir)
 
+    def test_pull_remote_main_image_urls_stops_after_open_and_closes_created_tab(self):
+        with self.make_temp_dir() as temp_dir:
+            cfg = self._config(temp_dir)
+            accounts.create_account("主店", "alias-a", debug_port=9222, config=cfg)
+            cdp = SimpleNamespace()
+            stop_values = iter([False, False, False, False, True])
+
+            with mock.patch("app.image_studio.chrome.is_running", return_value=True), \
+                    mock.patch(
+                        "app.image_studio.accounts.detect_login",
+                        return_value={"logged_in": True},
+                    ), \
+                    mock.patch(
+                        "app.image_studio.editor.open_product",
+                        return_value=cdp,
+                    ) as open_product, \
+                    mock.patch(
+                        "app.image_studio.editor.read_product_image_urls",
+                    ) as read_urls, \
+                    mock.patch(
+                        "app.image_studio.editor.close_readonly_product",
+                    ) as close_readonly:
+                with self.assertRaises(image_studio.ImageStudioPullCancelled):
+                    image_studio.pull_remote_main_image_urls(
+                        "alias-a",
+                        "51100639510",
+                        path=cfg["db_path"],
+                        config=cfg,
+                        should_stop=lambda: next(stop_values),
+                    )
+
+            open_product.assert_called_once()
+            read_urls.assert_not_called()
+            close_readonly.assert_called_once_with(cdp)
+
+        self.assert_removed(temp_dir)
+
+    def test_restore_original_asset_snapshot_restores_existing_state(self):
+        with self.make_temp_dir() as temp_dir:
+            db_path = os.path.join(temp_dir, "cmshopee.db")
+            db.init_db(db_path)
+            project = image_studio.create_or_get_project(
+                account_alias="alias-a",
+                account_slug="alias_a",
+                item_id="51100639510",
+                path=db_path,
+            )
+            assets = image_studio.sync_original_asset_urls(
+                project.id,
+                [
+                    {"index": 1, "src": "https://susercontent.com/one.jpg"},
+                    {"index": 2, "src": "https://susercontent.com/two.jpg"},
+                ],
+                path=db_path,
+            )
+            snapshot = [
+                {
+                    "id": assets[0].id,
+                    "status": image_studio.ASSET_STATUS_AVAILABLE,
+                    "source_order": 1,
+                },
+                {
+                    "id": assets[1].id,
+                    "status": image_studio.ASSET_STATUS_AVAILABLE,
+                    "source_order": 2,
+                },
+            ]
+            image_studio.sync_original_asset_urls(
+                project.id,
+                [{"index": 1, "src": "https://susercontent.com/two.jpg"}],
+                path=db_path,
+            )
+
+            restored = image_studio.restore_original_asset_snapshot(
+                project.id,
+                snapshot,
+                path=db_path,
+            )
+
+            self.assertEqual(
+                [
+                    (image_studio.ASSET_STATUS_AVAILABLE, 1),
+                    (image_studio.ASSET_STATUS_AVAILABLE, 2),
+                ],
+                [(asset.status, asset.source_order) for asset in restored],
+            )
+
+        self.assert_removed(temp_dir)
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_image_studio_images.py

@@ -29,6 +29,7 @@ class _Response:
         self.headers = dict(headers or {})
         self.url = url
         self.chunk_size = chunk_size
+        self.closed = False
 
     def iter_content(self, chunk_size=65536):
         if self.chunk_size:
@@ -37,6 +38,9 @@ class _Response:
             return
         yield self.content
 
+    def close(self):
+        self.closed = True
+
 
 class ImageStudioImageTests(TempDirMixin, unittest.TestCase):
     def _png_bytes(self, size=(20, 16), color=(80, 120, 200, 255)):
@@ -246,6 +250,33 @@ class ImageStudioImageTests(TempDirMixin, unittest.TestCase):
 
         self.assert_removed(temp_dir)
 
+    def test_download_remote_image_cancels_between_chunks_and_closes_response(self):
+        png = self._png_bytes(size=(120, 120))
+        response = _Response(
+            content=png,
+            headers={"Content-Type": "image/png"},
+            chunk_size=32,
+        )
+        session = SimpleNamespace(get=lambda *args, **kwargs: response)
+        checks = {"count": 0}
+
+        def should_stop():
+            checks["count"] += 1
+            return checks["count"] >= 4
+
+        with mock.patch(
+            "app.image_studio_images.socket.getaddrinfo",
+            return_value=self._public_dns(),
+        ):
+            with self.assertRaises(image_studio_images.ImageStudioImageCancelled):
+                image_studio_images.download_remote_image(
+                    "https://cdn.example.com/a.png",
+                    session=session,
+                    should_stop=should_stop,
+                )
+
+        self.assertTrue(response.closed)
+
     def test_import_original_files_copies_valid_images_deduplicates_and_limits(self):
         with self.make_temp_dir() as temp_dir:
             db_path = os.path.join(temp_dir, "cmshopee.db")

tests/test_product_suite_gui.py

@@ -1866,6 +1866,298 @@ class ProductSuiteGuiTests(TempDirMixin, unittest.TestCase):
 
         self.assert_removed(temp_dir)
 
+    def test_pull_button_stops_with_confirmation_and_repeated_click_is_ignored(self):
+        with self.make_temp_dir() as temp_dir:
+            config = self._config(temp_dir)
+            accounts.create_account("主店", "alias-a", debug_port=9222, config=config)
+            tab = ProductSuiteTab(config=config, db_path=config["db_path"])
+            self.addCleanup(tab.close)
+            state = tab._displayed_state
+            tab.item_id_edit.setText("51100639510")
+
+            with mock.patch.object(tab, "_confirm", return_value=True), \
+                    mock.patch.object(tab, "_start_thread", return_value=mock.Mock()):
+                tab.pull_main_images()
+
+            self.assertTrue(state.pull_running())
+            self.assertEqual("停止拉取蝦皮", tab.pull_button.text())
+            self.assertTrue(tab.pull_button.isEnabled())
+            worker = state.pull_worker
+
+            with mock.patch.object(tab, "_pull_stop_action", return_value="continue"):
+                tab.pull_main_images()
+            self.assertFalse(worker.is_cancelled())
+            self.assertFalse(state.pull_stop_requested)
+
+            with mock.patch.object(tab, "_pull_stop_action", return_value="keep"):
+                tab.pull_main_images()
+            self.assertTrue(worker.is_cancelled())
+            self.assertTrue(state.pull_stop_requested)
+            self.assertEqual("正在停止...", tab.pull_button.text())
+
+            stop_action = mock.Mock(return_value="clear_current")
+            with mock.patch.object(tab, "_pull_stop_action", stop_action):
+                tab.pull_main_images()
+            stop_action.assert_not_called()
+
+            state.pull_worker = None
+            state.pull_thread = None
+            tab._pull_run_states.clear()
+            state.pull_run_token = ""
+
+        self.assert_removed(temp_dir)
+
+    def test_stopping_pull_cancels_only_current_pull_downloads(self):
+        with self.make_temp_dir() as temp_dir:
+            config = self._config(temp_dir)
+            accounts.create_account("主店", "alias-a", debug_port=9222, config=config)
+            tab = ProductSuiteTab(config=config, db_path=config["db_path"])
+            self.addCleanup(tab.close)
+            state = tab._displayed_state
+            pull_worker = mock.Mock()
+            manual_worker = mock.Mock()
+            state.download_queue = [11, 12]
+            state.downloads = {
+                13: (pull_worker, mock.Mock()),
+                14: (manual_worker, mock.Mock()),
+            }
+            state.download_tokens = {
+                11: "pull-token",
+                12: "",
+                13: "pull-token",
+                14: "",
+            }
+            state.pull_download_asset_ids = {11, 13}
+
+            tab._cancel_pull_downloads(state)
+
+            self.assertEqual([12], state.download_queue)
+            self.assertEqual({13}, state.pull_download_asset_ids)
+            self.assertNotIn(11, state.download_tokens)
+            pull_worker.cancel.assert_called_once()
+            manual_worker.cancel.assert_not_called()
+
+        self.assert_removed(temp_dir)
+
+    def test_clear_stopped_pull_removes_only_new_unreferenced_remote_assets(self):
+        with self.make_temp_dir() as temp_dir:
+            config = self._config(temp_dir)
+            account = accounts.create_account(
+                "主店",
+                "alias-a",
+                debug_port=9222,
+                config=config,
+            )
+            project = image_studio.create_or_get_project(
+                account,
+                item_id="51100639510",
+                path=config["db_path"],
+            )
+            existing_remote = image_studio.sync_original_asset_urls(
+                project.id,
+                [{"index": 1, "src": "https://susercontent.com/existing.jpg"}],
+                path=config["db_path"],
+            )[0]
+            local_path = os.path.join(temp_dir, "local.png")
+            self._write_image(local_path)
+            local_asset = image_studio.add_asset(
+                project.id,
+                image_studio.ASSET_KIND_ORIGINAL,
+                local_path=local_path,
+                source_order=2,
+                path=config["db_path"],
+            )
+            before_assets = image_studio.list_assets(
+                project.id,
+                kind=image_studio.ASSET_KIND_ORIGINAL,
+                path=config["db_path"],
+            )
+            synced = image_studio.sync_original_asset_urls(
+                project.id,
+                [
+                    {"index": 1, "src": "https://susercontent.com/new-one.jpg"},
+                    {"index": 2, "src": "https://susercontent.com/new-two.jpg"},
+                ],
+                path=config["db_path"],
+            )
+            new_assets = [
+                asset
+                for asset in synced
+                if asset.remote_url
+                and asset.remote_url.endswith(("new-one.jpg", "new-two.jpg"))
+            ]
+            referenced = next(
+                asset for asset in new_assets if asset.remote_url.endswith("new-two.jpg")
+            )
+            image_studio.create_job(
+                project.id,
+                source_asset_id=referenced.id,
+                path=config["db_path"],
+            )
+
+            tab = ProductSuiteTab(config=config, db_path=config["db_path"])
+            self.addCleanup(tab.close)
+            state = tab._displayed_state
+            state.account_alias = "alias-a"
+            state.item_id = project.item_id
+            state.project_id = project.id
+            state.project_binding_state = project.binding_state
+            state.pull_run_token = "clear-pull"
+            state.pull_stop_requested = True
+            state.pull_cleanup_mode = "clear_current"
+            state.pull_before_asset_ids = {int(asset.id) for asset in before_assets}
+            state.pull_before_asset_states = [
+                {
+                    "id": int(asset.id),
+                    "status": asset.status,
+                    "source_order": int(asset.source_order),
+                }
+                for asset in before_assets
+            ]
+            state.pull_asset_ids = {int(asset.id) for asset in new_assets}
+            state.pull_started_at = time.monotonic()
+            tab._pull_run_states["clear-pull"] = state.key
+            messages = []
+            tab._message = lambda title, message, **kwargs: messages.append(
+                (title, message)
+            )
+
+            self.assertTrue(tab._finalize_pull(state, "clear-pull"))
+
+            remaining = {
+                asset.id: asset
+                for asset in image_studio.list_assets(
+                    project.id,
+                    kind=image_studio.ASSET_KIND_ORIGINAL,
+                    path=config["db_path"],
+                )
+            }
+            removed = next(
+                asset for asset in new_assets if asset.id != referenced.id
+            )
+            self.assertNotIn(removed.id, remaining)
+            self.assertIn(referenced.id, remaining)
+            self.assertEqual(
+                image_studio.ASSET_STATUS_AVAILABLE,
+                remaining[existing_remote.id].status,
+            )
+            self.assertEqual(1, remaining[existing_remote.id].source_order)
+            self.assertIn(local_asset.id, remaining)
+            self.assertGreater(
+                remaining[referenced.id].source_order,
+                remaining[local_asset.id].source_order,
+            )
+            self.assertEqual("拉取蝦皮主图已停止", messages[-1][0])
+            self.assertIn("清理1张", messages[-1][1])
+            self.assertIn("因引用保留1张", messages[-1][1])
+
+        self.assert_removed(temp_dir)
+
+    def test_old_pull_result_does_not_override_current_run(self):
+        with self.make_temp_dir() as temp_dir:
+            config = self._config(temp_dir)
+            accounts.create_account("主店", "alias-a", debug_port=9222, config=config)
+            tab = ProductSuiteTab(config=config, db_path=config["db_path"])
+            self.addCleanup(tab.close)
+            state = tab._displayed_state
+            state.pull_run_token = "current-pull"
+            state.pull_worker = mock.Mock()
+            tab._pull_run_states["current-pull"] = state.key
+
+            tab._on_pull_finished(
+                "old-pull",
+                {"project": None, "assets": [], "count": 0},
+            )
+
+            self.assertEqual("current-pull", state.pull_run_token)
+            self.assertIsNotNone(state.pull_worker)
+
+            state.pull_worker = None
+            state.pull_run_token = ""
+            tab._pull_run_states.clear()
+
+        self.assert_removed(temp_dir)
+
+    def test_pull_thread_finished_fallback_finalizes_requested_stop(self):
+        with self.make_temp_dir() as temp_dir:
+            config = self._config(temp_dir)
+            accounts.create_account("主店", "alias-a", debug_port=9222, config=config)
+            tab = ProductSuiteTab(config=config, db_path=config["db_path"])
+            self.addCleanup(tab.close)
+            state = tab._displayed_state
+            state.pull_run_token = "pull-fallback"
+            state.pull_stop_requested = True
+            state.pull_worker = mock.Mock()
+            state.pull_thread = mock.Mock()
+            state.pull_started_at = time.monotonic()
+            tab._pull_run_states["pull-fallback"] = state.key
+            messages = []
+            tab._message = lambda title, message, **kwargs: messages.append(
+                (title, message)
+            )
+
+            tab._handle_pull_thread_finished("pull-fallback")
+
+            self.assertEqual("", state.pull_run_token)
+            self.assertIsNone(state.pull_worker)
+            self.assertEqual("拉取蝦皮主图已停止", messages[-1][0])
+            self.assertEqual("拉取蝦皮主图", tab.pull_button.text())
+
+        self.assert_removed(temp_dir)
+
+    def test_pull_real_qthread_completion_restores_button_once(self):
+        with self.make_temp_dir() as temp_dir:
+            config = self._config(temp_dir)
+            account = accounts.create_account(
+                "主店",
+                "alias-a",
+                debug_port=9222,
+                config=config,
+            )
+            project = image_studio.create_or_get_project(
+                account,
+                item_id="51100639510",
+                path=config["db_path"],
+            )
+            tab = ProductSuiteTab(config=config, db_path=config["db_path"])
+            self.addCleanup(tab.close)
+            state = tab._displayed_state
+            tab.item_id_edit.setText(project.item_id)
+
+            with mock.patch.object(tab, "_confirm", return_value=True), \
+                    mock.patch(
+                        "app.gui.workers.image_studio.pull_remote_main_image_urls",
+                        return_value={
+                            "project": project,
+                            "assets": [],
+                            "account": account,
+                        },
+                    ):
+                tab.pull_main_images()
+                pull_thread = state.pull_thread
+                deadline = time.monotonic() + 3
+                while state.pull_running() and time.monotonic() < deadline:
+                    QTest.qWait(20)
+                    self.app.processEvents()
+                while pull_thread is not None and time.monotonic() < deadline:
+                    try:
+                        running = pull_thread.isRunning()
+                    except RuntimeError:
+                        pull_thread = None
+                        break
+                    if not running:
+                        break
+                    QTest.qWait(20)
+                    self.app.processEvents()
+
+            self.assertFalse(state.pull_running())
+            self.assertIsNone(state.pull_worker)
+            self.assertEqual("拉取蝦皮主图", tab.pull_button.text())
+            if pull_thread is not None:
+                self.assertFalse(pull_thread.isRunning())
+
+        self.assert_removed(temp_dir)
+
     def test_generation_confirmation_prevents_job_creation_and_retry_bypasses_it(self):
         with self.make_temp_dir() as temp_dir:
             config = self._config(temp_dir)

tests/test_workers.py

@@ -7,7 +7,7 @@ os.environ.setdefault("QT_QPA_PLATFORM", "offscreen")
 
 from _helpers import REPO_ROOT  # noqa: F401
 
-from app import workers
+from app import image_studio_images, workers
 
 if workers.QT_IMPORT_ERROR is not None:
     raise unittest.SkipTest("PySide6 未安装")
@@ -16,7 +16,10 @@ from PySide6.QtCore import QEventLoop, QTimer
 from PySide6.QtWidgets import QApplication
 
 from app.workers import BaseWorker, run_worker
-from app.gui.workers import ImageStudioDownloadOriginalWorker
+from app.gui.workers import (
+    ImageStudioDownloadOriginalWorker,
+    ImageStudioPullImagesWorker,
+)
 
 
 class DemoWorker(BaseWorker):
@@ -165,6 +168,34 @@ class WorkerTests(unittest.TestCase):
         self.assertEqual("蝦皮原主图下载失败，请稍后再次点击图片重试。", summary["error"])
         self.assertNotIn("https://", summary["error"])
 
+    def test_image_studio_pull_worker_returns_token_when_cancelled(self):
+        worker = ImageStudioPullImagesWorker(
+            "alias-a",
+            "51100639510",
+            pull_run_token="pull-token",
+        )
+        worker.cancel()
+
+        with mock.patch(
+            "app.gui.workers.image_studio.pull_remote_main_image_urls",
+        ) as pull:
+            summary = worker.execute()
+
+        pull.assert_not_called()
+        self.assertTrue(summary["cancelled"])
+        self.assertEqual("pull-token", summary["pull_run_token"])
+
+    def test_image_studio_original_download_maps_safe_cancel(self):
+        worker = ImageStudioDownloadOriginalWorker(12)
+
+        with mock.patch(
+            "app.gui.workers.image_studio_images.download_original_asset",
+            side_effect=image_studio_images.ImageStudioImageCancelled("停止"),
+        ):
+            summary = worker.execute()
+
+        self.assertEqual({"asset_id": 12, "cancelled": True}, summary)
+
 
 if __name__ == "__main__":
     unittest.main()