import os import sys import unittest from types import SimpleNamespace from unittest import mock sys.path.insert(0, os.path.dirname(__file__)) from _helpers import TempDirMixin from app import accounts, db, editor, image_studio class ImageStudioTests(TempDirMixin, unittest.TestCase): def _config(self, temp_dir): return { "data_dir": temp_dir, "db_path": os.path.join(temp_dir, "cmshopee.db"), "user_data_root": os.path.join(temp_dir, "chrome_user_data_dir"), "chrome_path": "chrome.exe", "default_debug_port": 9222, "debug_port_range": [9222, 9230], "cdp_ready_timeout": 1, } def test_init_db_adds_image_studio_tables_without_breaking_existing_tables(self): with self.make_temp_dir() as temp_dir: db_path = os.path.join(temp_dir, "cmshopee.db") db.init_db(db_path) db.init_db(db_path) conn = db.connect(db_path) try: tables = { row["name"] for row in conn.execute( "SELECT name FROM sqlite_master WHERE type = 'table'" ).fetchall() } self.assertTrue({"batches", "accounts", "tasks"}.issubset(tables)) self.assertTrue( { "image_studio_projects", "image_studio_assets", "image_studio_jobs", "image_studio_selections", }.issubset(tables) ) projects_columns = { row["name"] for row in conn.execute("PRAGMA table_info(image_studio_projects)").fetchall() } self.assertTrue( { "account_alias", "account_slug", "item_id", "storage_key", "binding_state", "target_main_count", "target_detail_count", "suite_settings_json", "deleted_at", }.issubset(projects_columns) ) jobs_columns = { row["name"] for row in conn.execute("PRAGMA table_info(image_studio_jobs)").fetchall() } self.assertIn("recovery_action", jobs_columns) finally: conn.close() self.assert_removed(temp_dir) def test_init_db_migrates_legacy_projects_with_default_suite_settings(self): with self.make_temp_dir() as temp_dir: db_path = os.path.join(temp_dir, "legacy-project.db") conn = db.connect(db_path) try: conn.execute( """ CREATE TABLE image_studio_projects ( id INTEGER PRIMARY KEY, account_alias TEXT NOT NULL, account_slug TEXT NOT NULL, account_name TEXT, item_id TEXT NOT NULL, target_main_count INTEGER NOT NULL DEFAULT 9, target_detail_count INTEGER NOT NULL DEFAULT 12, draft_prompt TEXT, status TEXT NOT NULL DEFAULT 'active', created_at TEXT NOT NULL, updated_at TEXT NOT NULL, deleted_at TEXT, deleted_reason TEXT, UNIQUE(account_alias, item_id) ) """ ) conn.execute( """ INSERT INTO image_studio_projects (id, account_alias, account_slug, item_id, created_at, updated_at) VALUES (1, 'alias', 'alias_slug', '51100639510', '2026-07-14', '2026-07-14') """ ) conn.commit() finally: conn.close() db.init_db(db_path) project = image_studio.get_project(1, path=db_path) self.assertEqual("{}", project.suite_settings_json) self.assertEqual({}, image_studio.project_suite_settings(project)) self.assertEqual("51100639510", project.storage_key) self.assertEqual(image_studio.PROJECT_BINDING_BOUND, project.binding_state) self.assert_removed(temp_dir) def test_project_crud_unique_per_account_and_image_dirs(self): with self.make_temp_dir() as temp_dir: db_path = os.path.join(temp_dir, "cmshopee.db") db.init_db(db_path) account = SimpleNamespace( alias="alias A", account_name="店铺A", slug="alias_a_slug", ) project = image_studio.create_or_get_project( account, item_id="51100639510", draft_prompt="初始提示词", path=db_path, ) same_project = image_studio.create_or_get_project( account_alias="alias A", account_slug="ignored_slug", item_id="51100639510", path=db_path, ) other_account_project = image_studio.create_or_get_project( account_alias="alias B", account_slug="alias_b_slug", item_id="51100639510", path=db_path, ) self.assertEqual(project.id, same_project.id) self.assertNotEqual(project.id, other_account_project.id) self.assertEqual("alias A", project.account_alias) self.assertEqual("alias_a_slug", project.account_slug) self.assertEqual("店铺A", project.account_name) self.assertEqual("初始提示词", project.draft_prompt) self.assertEqual("51100639510", project.storage_key) self.assertEqual(image_studio.PROJECT_BINDING_BOUND, project.binding_state) updated = image_studio.update_project_prompt(project.id, "二次提示词", path=db_path) self.assertEqual("二次提示词", updated.draft_prompt) suite_updated = image_studio.update_project_suite_settings( project.id, {"ratio": "3:4", "categories": {"白底图": 1}}, path=db_path, ) self.assertEqual( {"ratio": "3:4", "categories": {"白底图": 1}}, image_studio.project_suite_settings(suite_updated), ) dirs = image_studio.project_image_dirs(os.path.join(temp_dir, "images"), project) self.assertEqual( os.path.join( temp_dir, "images", "pool", "alias_a_slug", "51100639510", "originals", ), dirs["originals"], ) self.assertEqual( os.path.join( temp_dir, "images", "pool", "alias_a_slug", "51100639510", "generated", ), dirs["generated"], ) self.assertEqual( os.path.join( temp_dir, "images", "pool", "alias_a_slug", "51100639510", "exports", ), dirs["exports"], ) deleted = image_studio.soft_delete_project(project.id, "测试删除", path=db_path) self.assertEqual("测试删除", deleted.deleted_reason) self.assertEqual(1, len(image_studio.list_projects(path=db_path))) restored = image_studio.create_or_get_project( account, item_id="51100639510", path=db_path, ) self.assertEqual(project.id, restored.id) self.assertIsNone(restored.deleted_at) self.assert_removed(temp_dir) def test_draft_projects_bind_in_place_recover_and_conflict_safely(self): with self.make_temp_dir() as temp_dir: db_path = os.path.join(temp_dir, "cmshopee.db") db.init_db(db_path) account = SimpleNamespace( alias="alias-a", account_name="主店", slug="alias_a", ) draft = image_studio.create_draft_project( account, draft_prompt="临时草稿提示词", path=db_path, ) self.assertTrue(draft.item_id.startswith(image_studio.TEMPORARY_ITEM_PREFIX)) self.assertEqual(draft.item_id, draft.storage_key) self.assertEqual(image_studio.PROJECT_BINDING_DRAFT, draft.binding_state) self.assertTrue(image_studio.project_has_content(draft.id, path=db_path)) self.assertEqual( [draft.id], [ project.id for project in image_studio.list_recoverable_draft_projects( path=db_path ) ], ) image_studio.update_project_prompt(draft.id, "", path=db_path) self.assertFalse(image_studio.project_has_content(draft.id, path=db_path)) self.assertEqual([], image_studio.list_recoverable_draft_projects(path=db_path)) before_dirs = image_studio.project_image_dirs(os.path.join(temp_dir, "images"), draft) original = image_studio.add_asset( draft.id, image_studio.ASSET_KIND_ORIGINAL, local_path=os.path.join(temp_dir, "draft.png"), path=db_path, ) job = image_studio.create_job( draft.id, source_asset_id=original.id, path=db_path, ) image_studio.replace_selections(draft.id, "main", [original.id], path=db_path) self.assertTrue(image_studio.project_has_content(draft.id, path=db_path)) self.assertEqual([draft.id], [project.id for project in image_studio.list_recoverable_draft_projects(path=db_path)]) conflict = image_studio.create_or_get_project( account, item_id="51100639510", path=db_path, ) image_studio.soft_delete_project(conflict.id, "历史项目", path=db_path) with self.assertRaises(image_studio.ImageStudioProjectConflictError): image_studio.bind_draft_project(draft.id, "51100639510", path=db_path) bound = image_studio.bind_draft_project(draft.id, "51100639511", path=db_path) self.assertEqual(draft.id, bound.id) self.assertEqual("51100639511", bound.item_id) self.assertEqual(image_studio.PROJECT_BINDING_BOUND, bound.binding_state) self.assertEqual(draft.storage_key, bound.storage_key) self.assertEqual(before_dirs, image_studio.project_image_dirs(os.path.join(temp_dir, "images"), bound)) self.assertEqual(job.id, image_studio.list_jobs(bound.id, path=db_path)[0].id) self.assertEqual(original.id, image_studio.list_selections(bound.id, "main", path=db_path)[0].asset_id) self.assertEqual(bound, image_studio.bind_draft_project(bound.id, "51100639511", path=db_path)) empty = image_studio.create_draft_project(account, path=db_path) discarded = image_studio.discard_empty_draft_project(empty.id, path=db_path) self.assertIsNotNone(discarded.deleted_at) self.assert_removed(temp_dir) def test_asset_crud_parent_status_and_sorting(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", account_slug="alias_slug", item_id="51100639510", path=db_path, ) original_path = os.path.join(temp_dir, "old.png") generated_path = os.path.join(temp_dir, "generated.png") original = image_studio.add_asset( project.id, "original", remote_url="https://example.test/old.png", local_path=original_path, aspect_ratio="1:1", source_order=2, path=db_path, ) generated = image_studio.add_asset( project.id, "generated_main", local_path=generated_path, parent_asset_id=original.id, prompt="生成提示词", source_order=1, path=db_path, ) self.assertEqual(os.path.abspath(original_path), original.local_path) self.assertEqual(original.id, generated.parent_asset_id) self.assertEqual("生成提示词", generated.prompt) assets = image_studio.list_assets(project.id, path=db_path) self.assertEqual([generated.id, original.id], [asset.id for asset in assets]) missing = image_studio.mark_asset_status( original.id, image_studio.ASSET_STATUS_MISSING, path=db_path, ) self.assertEqual(image_studio.ASSET_STATUS_MISSING, missing.status) with self.assertRaises(db.DbError): image_studio.mark_asset_status(original.id, "deleted", path=db_path) available_assets = image_studio.list_assets( project.id, include_missing=False, path=db_path, ) self.assertEqual([generated.id], [asset.id for asset in available_assets]) self.assert_removed(temp_dir) def test_reorder_original_assets_requires_complete_project_order(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", account_slug="alias_slug", item_id="51100639510", path=db_path, ) first = image_studio.add_asset( project.id, image_studio.ASSET_KIND_ORIGINAL, local_path=os.path.join(temp_dir, "first.png"), source_order=1, path=db_path, ) second = image_studio.add_asset( project.id, image_studio.ASSET_KIND_ORIGINAL, local_path=os.path.join(temp_dir, "second.png"), source_order=2, path=db_path, ) reordered = image_studio.reorder_original_assets( project.id, [second.id, first.id], path=db_path, ) self.assertEqual([second.id, first.id], [asset.id for asset in reordered]) with self.assertRaisesRegex(db.DbError, "全部原图"): image_studio.reorder_original_assets(project.id, [first.id], path=db_path) self.assert_removed(temp_dir) def test_remove_asset_only_when_not_referenced(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", account_slug="alias_slug", item_id="51100639510", path=db_path, ) free_asset = image_studio.add_asset(project.id, "original", path=db_path) referenced_asset = image_studio.add_asset(project.id, "original", path=db_path) image_studio.create_job( project.id, source_asset_id=referenced_asset.id, path=db_path, ) counts = image_studio.asset_reference_counts(referenced_asset.id, path=db_path) self.assertEqual(1, counts["source_job"]) with self.assertRaisesRegex(db.DbError, "引用"): image_studio.remove_asset_if_unused(referenced_asset.id, path=db_path) removed = image_studio.remove_asset_if_unused(free_asset.id, path=db_path) self.assertEqual(free_asset.id, removed.id) self.assertIsNone(image_studio.get_asset(free_asset.id, path=db_path)) self.assert_removed(temp_dir) def test_remove_original_assets_is_atomic_and_reorders_remaining_assets(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", account_slug="alias_slug", item_id="51100639510", path=db_path, ) local_paths = [] assets = [] for index in range(1, 5): local_path = os.path.join(temp_dir, "original-%d.png" % index) with open(local_path, "wb") as fh: fh.write(b"image-%d" % index) local_paths.append(local_path) assets.append( image_studio.add_asset( project.id, image_studio.ASSET_KIND_ORIGINAL, local_path=local_path, source_order=index * 10, path=db_path, ) ) removed = image_studio.remove_original_assets_if_unused( project.id, [assets[0].id, assets[2].id, assets[0].id], path=db_path, ) self.assertEqual([assets[0].id, assets[2].id], [asset.id for asset in removed]) remaining = image_studio.list_assets( project.id, kind=image_studio.ASSET_KIND_ORIGINAL, path=db_path, ) self.assertEqual([assets[1].id, assets[3].id], [asset.id for asset in remaining]) self.assertEqual([1, 2], [asset.source_order for asset in remaining]) self.assertTrue(all(os.path.isfile(path) for path in local_paths)) self.assert_removed(temp_dir) def test_remove_original_assets_rejects_invalid_or_referenced_batch_without_partial_delete(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", account_slug="alias_slug", item_id="51100639510", path=db_path, ) other_project = image_studio.create_or_get_project( account_alias="alias", account_slug="alias_slug", item_id="51100639511", path=db_path, ) free_asset = image_studio.add_asset( project.id, image_studio.ASSET_KIND_ORIGINAL, source_order=1, path=db_path, ) referenced_asset = image_studio.add_asset( project.id, image_studio.ASSET_KIND_ORIGINAL, source_order=2, path=db_path, ) selected_asset = image_studio.add_asset( project.id, image_studio.ASSET_KIND_ORIGINAL, source_order=3, path=db_path, ) generated_asset = image_studio.add_asset( project.id, "generated", path=db_path, ) foreign_asset = image_studio.add_asset( other_project.id, image_studio.ASSET_KIND_ORIGINAL, path=db_path, ) image_studio.create_job( project.id, source_asset_id=referenced_asset.id, path=db_path, ) image_studio.replace_selections( project.id, "main", [selected_asset.id], path=db_path, ) with self.assertRaisesRegex(db.DbError, "引用"): image_studio.remove_original_assets_if_unused( project.id, [free_asset.id, referenced_asset.id], path=db_path, ) self.assertIsNotNone(image_studio.get_asset(free_asset.id, path=db_path)) self.assertIsNotNone(image_studio.get_asset(referenced_asset.id, path=db_path)) with self.assertRaisesRegex(db.DbError, "引用"): image_studio.remove_original_assets_if_unused( project.id, [free_asset.id, selected_asset.id], path=db_path, ) self.assertIsNotNone(image_studio.get_asset(free_asset.id, path=db_path)) self.assertIsNotNone(image_studio.get_asset(selected_asset.id, path=db_path)) for invalid_id in (generated_asset.id, foreign_asset.id, 999999): with self.assertRaisesRegex(db.DbError, "不属于当前项目"): image_studio.remove_original_assets_if_unused( project.id, [free_asset.id, invalid_id], path=db_path, ) self.assertIsNotNone(image_studio.get_asset(free_asset.id, path=db_path)) with self.assertRaisesRegex(db.DbError, "请选择"): image_studio.remove_original_assets_if_unused(project.id, [], path=db_path) self.assert_removed(temp_dir) def test_sync_original_asset_urls_is_idempotent_and_marks_missing(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", account_slug="alias_slug", item_id="51100639510", path=db_path, ) first_sync = image_studio.sync_original_asset_urls( project.id, [ {"index": 1, "src": "https://susercontent.com/a.jpg"}, {"index": 2, "src": "https://susercontent.com/b.jpg"}, ], path=db_path, ) second_sync = image_studio.sync_original_asset_urls( project.id, [ {"index": 1, "src": "https://susercontent.com/a.jpg"}, {"index": 2, "src": "https://susercontent.com/b.jpg"}, ], path=db_path, ) third_sync = image_studio.sync_original_asset_urls( project.id, [{"index": 1, "src": "https://susercontent.com/b.jpg"}], path=db_path, ) self.assertEqual([asset.id for asset in first_sync], [asset.id for asset in second_sync]) self.assertEqual(2, len(third_sync)) by_url = {asset.remote_url: asset for asset in third_sync} self.assertEqual(image_studio.ASSET_STATUS_MISSING, by_url["https://susercontent.com/a.jpg"].status) self.assertEqual(image_studio.ASSET_STATUS_AVAILABLE, by_url["https://susercontent.com/b.jpg"].status) self.assertEqual(1, by_url["https://susercontent.com/b.jpg"].source_order) self.assertIsNone(by_url["https://susercontent.com/b.jpg"].local_path) self.assert_removed(temp_dir) def test_sync_original_asset_urls_preserves_local_upload_and_caps_active_assets(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", account_slug="alias_slug", item_id="51100639510", path=db_path, ) local = image_studio.add_asset( project.id, image_studio.ASSET_KIND_ORIGINAL, local_path=os.path.join(temp_dir, "local.png"), source_order=1, path=db_path, ) image_studio.sync_original_asset_urls( project.id, [ {"index": index, "src": "https://susercontent.com/%d.jpg" % index} for index in range(1, 8) ], max_assets=3, path=db_path, ) active = image_studio.list_assets( project.id, kind=image_studio.ASSET_KIND_ORIGINAL, include_missing=False, path=db_path, ) self.assertEqual(3, len(active)) self.assertIn(local.id, [asset.id for asset in active]) self.assertEqual( image_studio.ASSET_STATUS_AVAILABLE, image_studio.get_asset(local.id, path=db_path).status, ) image_studio.sync_original_asset_urls(project.id, [], max_assets=3, path=db_path) self.assertEqual( image_studio.ASSET_STATUS_AVAILABLE, image_studio.get_asset(local.id, path=db_path).status, ) self.assert_removed(temp_dir) def test_job_lifecycle_and_resumable_query(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", account_slug="alias_slug", item_id="51100639510", path=db_path, ) source_asset = image_studio.add_asset(project.id, "original", path=db_path) job = image_studio.create_job( project.id, source_asset_id=source_asset.id, job_type="main", task_key="stable-task-key", prompt="生成主图", path=db_path, ) self.assertEqual("stable-task-key", job.task_key) self.assertEqual("pending", job.status) self.assertEqual("cmhub", job.provider) self.assertEqual(image_studio.JOB_RECOVERY_REGENERATE, job.recovery_action) submitted = image_studio.set_job_submitted( job.id, "cmhub-task-1", call_id="call-1", points_cost=2, points_balance=98, path=db_path, ) self.assertEqual("submitted", submitted.status) self.assertEqual("cmhub-task-1", submitted.task_id) self.assertEqual(2, submitted.points_cost) self.assertEqual(image_studio.JOB_RECOVERY_RESUME, submitted.recovery_action) self.assertEqual([job.id], [item.id for item in image_studio.list_resumable_jobs(path=db_path)]) running = image_studio.update_job_status( job.id, "running", increment_attempts=True, path=db_path, ) self.assertEqual("running", running.status) self.assertEqual(1, running.attempts) with self.assertRaises(db.DbError): image_studio.update_job_status(job.id, "unknown", path=db_path) output_asset = image_studio.add_asset( project.id, "generated_main", parent_asset_id=source_asset.id, path=db_path, ) succeeded = image_studio.update_job_status( job.id, "succeeded", output_asset_id=output_asset.id, points_balance=96, path=db_path, ) self.assertEqual("succeeded", succeeded.status) self.assertEqual(output_asset.id, succeeded.output_asset_id) self.assertEqual(96, succeeded.points_balance) self.assertEqual(image_studio.JOB_RECOVERY_NONE, succeeded.recovery_action) self.assertIsNotNone(succeeded.finished_at) self.assertEqual([], image_studio.list_resumable_jobs(path=db_path)) terminal = image_studio.create_job( project.id, source_asset_id=source_asset.id, task_key="terminal-task-key", path=db_path, ) image_studio.set_job_submitted(terminal.id, "cmhub-terminal", path=db_path) terminal = image_studio.update_job_status( terminal.id, "failed", error="上游生成失败", recovery_action=image_studio.JOB_RECOVERY_REGENERATE, path=db_path, ) self.assertEqual(image_studio.JOB_RECOVERY_REGENERATE, terminal.recovery_action) self.assertEqual( [], image_studio.list_resumable_jobs(path=db_path, include_failed_downloads=True), ) self.assertEqual( [terminal.id, job.id], [item.id for item in image_studio.list_jobs(project.id, path=db_path)], ) self.assertEqual( [terminal.id], [ item.id for item in image_studio.list_jobs( project.id, statuses=["failed"], path=db_path, ) ], ) with self.assertRaisesRegex(db.DbError, "状态无效"): image_studio.list_jobs(project.id, statuses=["unknown"], path=db_path) with self.assertRaises(db.DbError): image_studio.create_job( project.id, task_key="stable-task-key", path=db_path, ) self.assert_removed(temp_dir) def test_init_db_migrates_legacy_image_studio_job_recovery_action(self): with self.make_temp_dir() as temp_dir: db_path = os.path.join(temp_dir, "legacy.db") conn = db.connect(db_path) try: conn.execute( """ CREATE TABLE image_studio_jobs ( id INTEGER PRIMARY KEY, project_id INTEGER NOT NULL, source_asset_id INTEGER, output_asset_id INTEGER, generation_source TEXT NOT NULL DEFAULT 'cmhub', provider TEXT NOT NULL DEFAULT 'cmhub', job_type TEXT NOT NULL, task_key TEXT NOT NULL UNIQUE, task_id TEXT, status TEXT NOT NULL DEFAULT 'pending', prompt TEXT, error TEXT, attempts INTEGER NOT NULL DEFAULT 0, points_cost INTEGER, points_balance INTEGER, call_id TEXT, created_at TEXT NOT NULL, updated_at TEXT NOT NULL, submitted_at TEXT, finished_at TEXT ) """ ) conn.executemany( """ INSERT INTO image_studio_jobs (id, project_id, job_type, task_key, task_id, status, created_at, updated_at) VALUES (?, 1, 'main', ?, ?, ?, '2026-07-13T00:00:00', '2026-07-13T00:00:00') """, [ (1, "legacy-submitted", "task-submitted", "submitted"), (2, "legacy-running", "task-running", "running"), (3, "legacy-failed", "task-failed", "failed"), (4, "legacy-success", "task-success", "succeeded"), ], ) conn.commit() finally: conn.close() db.init_db(db_path) conn = db.connect(db_path) try: recovery_actions = { row["task_key"]: row["recovery_action"] for row in conn.execute( "SELECT task_key, recovery_action FROM image_studio_jobs ORDER BY id" ).fetchall() } finally: conn.close() self.assertEqual(image_studio.JOB_RECOVERY_RESUME, recovery_actions["legacy-submitted"]) self.assertEqual(image_studio.JOB_RECOVERY_RESUME, recovery_actions["legacy-running"]) self.assertEqual(image_studio.JOB_RECOVERY_REGENERATE, recovery_actions["legacy-failed"]) self.assertEqual(image_studio.JOB_RECOVERY_NONE, recovery_actions["legacy-success"]) self.assert_removed(temp_dir) def test_selections_are_consecutive_unique_and_replaceable(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", account_slug="alias_slug", item_id="51100639510", path=db_path, ) other_project = image_studio.create_or_get_project( account_alias="other", account_slug="other_slug", item_id="51100639510", path=db_path, ) first = image_studio.add_asset(project.id, "generated_main", source_order=1, path=db_path) second = image_studio.add_asset(project.id, "generated_main", source_order=2, path=db_path) third = image_studio.add_asset(project.id, "generated_detail", source_order=3, path=db_path) other_asset = image_studio.add_asset(other_project.id, "generated_main", path=db_path) main = image_studio.replace_selections(project.id, "main", [second.id, first.id], path=db_path) self.assertEqual([1, 2], [selection.position for selection in main]) self.assertEqual([second.id, first.id], [selection.asset_id for selection in main]) detail = image_studio.replace_selections(project.id, "detail", [second.id, third.id], path=db_path) self.assertEqual([second.id, third.id], [selection.asset_id for selection in detail]) replaced = image_studio.replace_selections(project.id, "main", [first.id], path=db_path) self.assertEqual([1], [selection.position for selection in replaced]) self.assertEqual([first.id], [selection.asset_id for selection in replaced]) all_selections = image_studio.list_selections(project.id, path=db_path) self.assertEqual( [("detail", 1, second.id), ("detail", 2, third.id), ("main", 1, first.id)], [ (selection.selection_type, selection.position, selection.asset_id) for selection in all_selections ], ) with self.assertRaises(db.DbError): image_studio.replace_selections(project.id, "main", [first.id, first.id], path=db_path) with self.assertRaises(db.DbError): image_studio.replace_selections(project.id, "invalid", [first.id], path=db_path) with self.assertRaises(db.DbError): image_studio.replace_selections(project.id, "main", [other_asset.id], path=db_path) with self.assertRaises(db.DbError): image_studio.create_job(project.id, source_asset_id=other_asset.id, path=db_path) self.assert_removed(temp_dir) def test_pull_remote_main_image_urls_reuses_running_chrome_and_writes_snapshot(self): with self.make_temp_dir() as temp_dir: cfg = self._config(temp_dir) account = accounts.create_account("主店", "alias-a", debug_port=9222, config=cfg) images = [ {"index": index, "src": f"https://susercontent.com/main-{index}.jpg"} for index in range(1, 10) ] cdp = SimpleNamespace() with mock.patch("app.image_studio.chrome.is_running", return_value=True) as is_running, \ mock.patch("app.image_studio.accounts.launch_for_login") as launch_for_login, \ mock.patch( "app.image_studio.accounts.detect_login", return_value={"logged_in": True, "cookie_names": ["SPC_ST"]}, ) as detect_login, \ mock.patch("app.image_studio.editor.open_product", return_value=cdp) as open_product, \ mock.patch( "app.image_studio.editor.read_product_image_urls", return_value=images, ) as read_urls, \ mock.patch("app.image_studio.editor.close_readonly_product") as close_readonly, \ mock.patch("app.image_studio.editor.change_title") as change_title, \ mock.patch("app.image_studio.editor.replace_cover") as replace_cover, \ mock.patch("app.image_studio.editor.click_update") as click_update: result = image_studio.pull_remote_main_image_urls( "alias-a", "51100639510", path=cfg["db_path"], config=cfg, ) second = image_studio.pull_remote_main_image_urls( "alias-a", "51100639510", path=cfg["db_path"], config=cfg, ) self.assertEqual("alias-a", result["project"].account_alias) self.assertEqual("51100639510", result["project"].item_id) self.assertTrue(result["readiness"]["reused"]) self.assertFalse(result["readiness"]["login_uncertain"]) self.assertEqual(images, result["images"]) self.assertEqual([asset.id for asset in result["assets"]], [asset.id for asset in second["assets"]]) self.assertEqual(9, len(result["assets"])) self.assertEqual( list(range(1, 10)), [asset.source_order for asset in result["assets"]], ) self.assertTrue(all(asset.kind == image_studio.ASSET_KIND_ORIGINAL for asset in result["assets"])) self.assertTrue(all(asset.local_path is None for asset in result["assets"])) is_running.assert_called_with(9222) launch_for_login.assert_not_called() detect_login.assert_called() open_product.assert_called_with(account, "51100639510", on_step=None, bring_to_front=False) self.assertEqual(2, read_urls.call_count) self.assertEqual(2, close_readonly.call_count) change_title.assert_not_called() replace_cover.assert_not_called() click_update.assert_not_called() self.assert_removed(temp_dir) def test_pull_remote_main_image_urls_launches_chrome_when_needed(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() with mock.patch("app.image_studio.chrome.is_running", return_value=False) as is_running, \ mock.patch( "app.image_studio.accounts.launch_for_login", return_value={"launched": True, "reused": False, "debug_port": 9222}, ) as launch_for_login, \ mock.patch( "app.image_studio.accounts.detect_login", return_value={"logged_in": True}, ), \ mock.patch("app.image_studio.editor.open_product", return_value=cdp), \ mock.patch( "app.image_studio.editor.read_product_image_urls", return_value=[{"index": 1, "src": "https://susercontent.com/main.jpg"}], ), \ mock.patch("app.image_studio.editor.close_readonly_product"): result = image_studio.pull_remote_main_image_urls( "alias-a", "51100639510", path=cfg["db_path"], config=cfg, ) self.assertTrue(result["readiness"]["launched"]) is_running.assert_called_once_with(9222) launch_for_login.assert_called_once() self.assert_removed(temp_dir) def test_pull_remote_main_image_urls_blocks_definitive_logged_out(self): with self.make_temp_dir() as temp_dir: cfg = self._config(temp_dir) accounts.create_account("主店", "alias-a", debug_port=9222, config=cfg) with mock.patch("app.image_studio.chrome.is_running", return_value=True), \ mock.patch( "app.image_studio.accounts.detect_login", return_value={ "logged_in": False, "reason": "LOGIN_PAGE", "url": "https://accounts.shopee.tw/seller/login", "cookie_names": [], }, ), \ mock.patch("app.image_studio.editor.open_product") as open_product: with self.assertRaisesRegex(image_studio.ImageStudioError, "未登录"): image_studio.pull_remote_main_image_urls( "alias-a", "51100639510", path=cfg["db_path"], config=cfg, ) open_product.assert_not_called() self.assert_removed(temp_dir) def test_pull_remote_main_image_urls_continues_when_login_status_uncertain(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() with mock.patch("app.image_studio.chrome.is_running", return_value=True), \ mock.patch( "app.image_studio.accounts.detect_login", return_value={ "logged_in": False, "reason": "NO_SESSION_COOKIE", "url": "https://seller.shopee.tw/portal/", "cookie_names": [], }, ), \ mock.patch("app.image_studio.editor.open_product", return_value=cdp) as open_product, \ mock.patch( "app.image_studio.editor.read_product_image_urls", return_value=[{"index": 1, "src": "https://susercontent.com/main.jpg"}], ), \ mock.patch("app.image_studio.editor.close_readonly_product"): result = image_studio.pull_remote_main_image_urls( "alias-a", "51100639510", path=cfg["db_path"], config=cfg, ) self.assertTrue(result["readiness"]["login_uncertain"]) open_product.assert_called_once() self.assert_removed(temp_dir) def test_pull_remote_main_image_urls_preserves_product_unavailable_toast_error(self): with self.make_temp_dir() as temp_dir: cfg = self._config(temp_dir) accounts.create_account("主店", "alias-a", debug_port=9222, config=cfg) 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", side_effect=editor.EditorError("商品失效：please input correct product id"), ), \ mock.patch("app.image_studio.editor.close_readonly_product") as close_readonly: with self.assertRaisesRegex(image_studio.ImageStudioError, "please input correct product id"): image_studio.pull_remote_main_image_urls( "alias-a", "bad-item", path=cfg["db_path"], config=cfg, ) close_readonly.assert_not_called() 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()