feat(product-suite): make image pulls cancellable

2026-07-16 19:37:22 +08:00
parent c9afaad24b
commit 75af87d441
12 changed files with 1157 additions and 88 deletions
@@ -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()