cmshoppe/tests/test_image_studio.py

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",
                        "current_generation_round_key",
                        "deleted_at",
                    }.issubset(projects_columns)
                )
                jobs_columns = {
                    row["name"]
                    for row in conn.execute("PRAGMA table_info(image_studio_jobs)").fetchall()
                }
                self.assertTrue(
                    {
                        "recovery_action",
                        "generation_round_key",
                        "generation_slot_index",
                    }.issubset(jobs_columns)
                )
                indexes = {
                    row["name"]
                    for row in conn.execute(
                        "PRAGMA index_list(image_studio_jobs)"
                    ).fetchall()
                }
                self.assertIn("idx_image_studio_jobs_generation_round", indexes)
            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"])

            legacy_jobs = image_studio.list_generation_round_current_jobs(
                1,
                None,
                path=db_path,
            )
            self.assertEqual([1, 2, 3, 4], [job.id for job in legacy_jobs])
            self.assertTrue(
                all(job.generation_round_key is None for job in legacy_jobs)
            )

        self.assert_removed(temp_dir)

    def test_generation_round_queries_keep_slots_attempts_and_project_boundaries(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,
            )
            other_project = image_studio.create_or_get_project(
                account_alias="alias-b",
                account_slug="alias-b",
                item_id="51100639510",
                path=db_path,
            )
            source = image_studio.add_asset(
                project.id,
                image_studio.ASSET_KIND_ORIGINAL,
                path=db_path,
            )
            legacy = image_studio.create_job(
                project.id,
                source_asset_id=source.id,
                task_key="legacy-generation-round",
                path=db_path,
            )
            round_one = "round-one"
            first = image_studio.create_job(
                project.id,
                source_asset_id=source.id,
                task_key="round-one-slot-0",
                generation_round_key=round_one,
                generation_slot_index=0,
                path=db_path,
            )
            first = image_studio.update_job_status(
                first.id,
                "succeeded",
                path=db_path,
            )
            failed = image_studio.create_job(
                project.id,
                source_asset_id=source.id,
                task_key="round-one-slot-1-failed",
                generation_round_key=round_one,
                generation_slot_index=1,
                path=db_path,
            )
            failed = image_studio.update_job_status(
                failed.id,
                "failed",
                path=db_path,
            )
            retry = image_studio.create_job(
                project.id,
                source_asset_id=source.id,
                task_key="round-one-slot-1-retry",
                generation_round_key=round_one,
                generation_slot_index=1,
                path=db_path,
            )
            retry = image_studio.update_job_status(
                retry.id,
                "succeeded",
                path=db_path,
            )
            failed_round = "round-two-all-failed"
            failed_round_job = image_studio.create_job(
                project.id,
                source_asset_id=source.id,
                task_key="round-two-slot-0",
                generation_round_key=failed_round,
                generation_slot_index=0,
                path=db_path,
            )
            image_studio.update_job_status(
                failed_round_job.id,
                "failed",
                path=db_path,
            )

            current = image_studio.set_current_generation_round(
                project.id,
                round_one,
                path=db_path,
            )
            self.assertEqual(round_one, current.current_generation_round_key)
            self.assertEqual(
                round_one,
                image_studio.get_current_generation_round(project.id, path=db_path),
            )
            self.assertFalse(
                image_studio.promote_generation_round_if_success(
                    project.id,
                    failed_round,
                    path=db_path,
                )
            )
            self.assertEqual(
                round_one,
                image_studio.get_current_generation_round(project.id, path=db_path),
            )
            with self.assertRaisesRegex(db.DbError, "不属于"):
                image_studio.set_current_generation_round(
                    other_project.id,
                    round_one,
                    path=db_path,
                )

            current_jobs = image_studio.list_generation_round_current_jobs(
                project.id,
                round_one,
                path=db_path,
            )
            self.assertEqual([first.id, retry.id], [job.id for job in current_jobs])
            self.assertEqual([0, 1], [job.generation_slot_index for job in current_jobs])
            attempts = image_studio.list_generation_round_attempts(
                project.id,
                round_one,
                path=db_path,
            )
            self.assertEqual([first.id, failed.id, retry.id], [job.id for job in attempts])
            legacy_jobs = image_studio.list_generation_round_current_jobs(
                project.id,
                None,
                path=db_path,
            )
            self.assertEqual([legacy.id], [job.id for job in legacy_jobs])

            rounds = image_studio.list_generation_rounds(project.id, path=db_path)
            self.assertEqual(
                [failed_round, round_one, None],
                [round_.generation_round_key for round_ in rounds],
            )
            summary = rounds[1]
            self.assertTrue(summary.is_current)
            self.assertEqual(3, summary.job_count)
            self.assertEqual(2, summary.slot_count)
            self.assertEqual(1, summary.retry_count)
            self.assertEqual(2, summary.succeeded_count)
            self.assertEqual(1, summary.failed_count)
            self.assertTrue(rounds[-1].is_legacy)

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