1812 lines
74 KiB
Python
1812 lines
74 KiB
Python
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_formal_item_id_normalization_removes_format_characters_only(self):
|
|
zero_width_item_id = "51100639510" + chr(0x200B)
|
|
self.assertEqual(
|
|
"51100639510",
|
|
image_studio.normalize_item_id_input(" " + zero_width_item_id + " "),
|
|
)
|
|
self.assertTrue(image_studio.is_formal_item_id(zero_width_item_id))
|
|
self.assertFalse(image_studio.is_formal_item_id("51100639510-1"))
|
|
self.assertFalse(image_studio.is_formal_item_id("商品51100639510"))
|
|
self.assertFalse(image_studio.is_formal_item_id("１２３"))
|
|
|
|
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",
|
|
"reference_asset_ids",
|
|
"run_session_id",
|
|
}.issubset(jobs_columns)
|
|
)
|
|
assets_columns = {
|
|
row["name"]
|
|
for row in conn.execute("PRAGMA table_info(image_studio_assets)").fetchall()
|
|
}
|
|
self.assertTrue(
|
|
{
|
|
"requested_output_size",
|
|
"rendered_width",
|
|
"rendered_height",
|
|
}.issubset(assets_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_remove_original_assets_allows_completed_history_source_and_reference(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 = image_studio.add_asset(
|
|
project.id,
|
|
image_studio.ASSET_KIND_ORIGINAL,
|
|
source_order=1,
|
|
path=db_path,
|
|
)
|
|
reference = image_studio.add_asset(
|
|
project.id,
|
|
image_studio.ASSET_KIND_ORIGINAL,
|
|
source_order=2,
|
|
path=db_path,
|
|
)
|
|
output = image_studio.add_asset(
|
|
project.id,
|
|
"generated_main",
|
|
parent_asset_id=source.id,
|
|
path=db_path,
|
|
)
|
|
job = image_studio.create_job(
|
|
project.id,
|
|
source_asset_id=source.id,
|
|
reference_asset_ids=[reference.id],
|
|
path=db_path,
|
|
)
|
|
job = image_studio.update_job_status(
|
|
job.id,
|
|
"succeeded",
|
|
output_asset_id=output.id,
|
|
path=db_path,
|
|
)
|
|
self.assertEqual(image_studio.JOB_RECOVERY_NONE, job.recovery_action)
|
|
|
|
removed = image_studio.remove_original_assets_if_unused(
|
|
project.id,
|
|
[source.id, reference.id],
|
|
path=db_path,
|
|
)
|
|
|
|
self.assertEqual([source.id, reference.id], [asset.id for asset in removed])
|
|
self.assertIsNone(image_studio.get_asset(source.id, path=db_path))
|
|
self.assertIsNone(image_studio.get_asset(reference.id, path=db_path))
|
|
saved_job = image_studio.get_job(job.id, path=db_path)
|
|
saved_output = image_studio.get_asset(output.id, path=db_path)
|
|
self.assertIsNone(saved_job.source_asset_id)
|
|
self.assertEqual("[%d]" % reference.id, saved_job.reference_asset_ids)
|
|
self.assertIsNone(saved_output.parent_asset_id)
|
|
|
|
self.assert_removed(temp_dir)
|
|
|
|
def test_remove_original_assets_blocks_active_and_retryable_source_or_reference(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,
|
|
)
|
|
statuses = ("pending", "submitted", "running", "failed", "expired", "cancelled")
|
|
for index, status in enumerate(statuses, 1):
|
|
source = image_studio.add_asset(
|
|
project.id,
|
|
image_studio.ASSET_KIND_ORIGINAL,
|
|
source_order=index * 2 - 1,
|
|
path=db_path,
|
|
)
|
|
reference = image_studio.add_asset(
|
|
project.id,
|
|
image_studio.ASSET_KIND_ORIGINAL,
|
|
source_order=index * 2,
|
|
path=db_path,
|
|
)
|
|
job = image_studio.create_job(
|
|
project.id,
|
|
source_asset_id=source.id,
|
|
reference_asset_ids=[reference.id],
|
|
path=db_path,
|
|
)
|
|
if status in {"submitted", "running"}:
|
|
job = image_studio.set_job_submitted(
|
|
job.id,
|
|
"task-%d" % index,
|
|
path=db_path,
|
|
)
|
|
if status == "running":
|
|
job = image_studio.update_job_status(job.id, "running", path=db_path)
|
|
elif status in {"failed", "expired", "cancelled"}:
|
|
job = image_studio.update_job_status(
|
|
job.id,
|
|
status,
|
|
recovery_action=image_studio.JOB_RECOVERY_REGENERATE,
|
|
path=db_path,
|
|
)
|
|
self.assertEqual(status, job.status)
|
|
with self.assertRaisesRegex(db.DbError, "可继续处理的生成任务引用"):
|
|
image_studio.remove_original_assets_if_unused(
|
|
project.id,
|
|
[reference.id],
|
|
path=db_path,
|
|
)
|
|
self.assertIsNotNone(image_studio.get_asset(reference.id, path=db_path))
|
|
|
|
self.assert_removed(temp_dir)
|
|
|
|
def test_remove_original_assets_blocks_invalid_reference_snapshot(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 = image_studio.add_asset(
|
|
project.id,
|
|
image_studio.ASSET_KIND_ORIGINAL,
|
|
path=db_path,
|
|
)
|
|
job = image_studio.create_job(
|
|
project.id,
|
|
source_asset_id=source.id,
|
|
path=db_path,
|
|
)
|
|
job = image_studio.update_job_status(job.id, "succeeded", path=db_path)
|
|
conn = db.connect(db_path)
|
|
try:
|
|
with conn:
|
|
conn.execute(
|
|
"UPDATE image_studio_jobs SET reference_asset_ids = ? WHERE id = ?",
|
|
("{坏快照", job.id),
|
|
)
|
|
finally:
|
|
conn.close()
|
|
|
|
with self.assertRaisesRegex(db.DbError, "参考图快照无效"):
|
|
image_studio.remove_original_assets_if_unused(
|
|
project.id,
|
|
[source.id],
|
|
path=db_path,
|
|
)
|
|
self.assertIsNotNone(image_studio.get_asset(source.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.assertTrue(
|
|
all(job.reference_asset_ids is None for job in legacy_jobs)
|
|
)
|
|
|
|
self.assert_removed(temp_dir)
|
|
|
|
def test_create_job_freezes_valid_reference_asset_ids(self):
|
|
with self.make_temp_dir() as temp_dir:
|
|
config = self._config(temp_dir)
|
|
db.init_db(config["db_path"])
|
|
project = image_studio.create_or_get_project(
|
|
account_alias="店铺",
|
|
account_slug="shop",
|
|
item_id="51100639510",
|
|
path=config["db_path"],
|
|
)
|
|
source = image_studio.add_asset(
|
|
project.id,
|
|
image_studio.ASSET_KIND_ORIGINAL,
|
|
path=config["db_path"],
|
|
)
|
|
reference = image_studio.add_asset(
|
|
project.id,
|
|
image_studio.ASSET_KIND_ORIGINAL,
|
|
path=config["db_path"],
|
|
)
|
|
job = image_studio.create_job(
|
|
project.id,
|
|
source_asset_id=source.id,
|
|
reference_asset_ids=[reference.id],
|
|
path=config["db_path"],
|
|
)
|
|
|
|
self.assertEqual("[%d]" % reference.id, job.reference_asset_ids)
|
|
self.assertEqual([reference.id], image_studio.job_reference_asset_ids(job))
|
|
with self.assertRaisesRegex(db.DbError, "不能包含主图"):
|
|
image_studio.create_job(
|
|
project.id,
|
|
source_asset_id=source.id,
|
|
reference_asset_ids=[source.id],
|
|
path=config["db_path"],
|
|
)
|
|
|
|
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,
|
|
)
|
|
other_source = image_studio.add_asset(
|
|
other_project.id,
|
|
image_studio.ASSET_KIND_ORIGINAL,
|
|
path=db_path,
|
|
)
|
|
other_first = image_studio.create_job(
|
|
other_project.id,
|
|
source_asset_id=other_source.id,
|
|
task_key="other-first-generation-round",
|
|
generation_round_key="other-failed-round",
|
|
generation_slot_index=0,
|
|
path=db_path,
|
|
)
|
|
image_studio.update_job_status(
|
|
other_first.id,
|
|
"succeeded",
|
|
path=db_path,
|
|
)
|
|
other_failed = image_studio.create_job(
|
|
other_project.id,
|
|
source_asset_id=other_source.id,
|
|
task_key="other-failed-generation-round",
|
|
generation_round_key="other-failed-round",
|
|
generation_slot_index=0,
|
|
path=db_path,
|
|
)
|
|
image_studio.update_job_status(
|
|
other_failed.id,
|
|
"failed",
|
|
path=db_path,
|
|
)
|
|
legacy = image_studio.create_job(
|
|
project.id,
|
|
source_asset_id=source.id,
|
|
task_key="legacy-generation-round",
|
|
path=db_path,
|
|
)
|
|
legacy = image_studio.update_job_status(
|
|
legacy.id,
|
|
"succeeded",
|
|
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],
|
|
)
|
|
self.assertEqual(
|
|
[failed_round],
|
|
[
|
|
round_.generation_round_key
|
|
for round_ in image_studio.list_generation_rounds(
|
|
project.id,
|
|
limit=1,
|
|
offset=0,
|
|
path=db_path,
|
|
)
|
|
],
|
|
)
|
|
self.assertEqual(
|
|
[round_one],
|
|
[
|
|
round_.generation_round_key
|
|
for round_ in image_studio.list_generation_rounds(
|
|
project.id,
|
|
limit=1,
|
|
offset=1,
|
|
path=db_path,
|
|
)
|
|
],
|
|
)
|
|
self.assertEqual(
|
|
[None],
|
|
[
|
|
round_.generation_round_key
|
|
for round_ in image_studio.list_generation_rounds(
|
|
project.id,
|
|
limit=1,
|
|
offset=2,
|
|
path=db_path,
|
|
)
|
|
],
|
|
)
|
|
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)
|
|
|
|
success_summary = image_studio.get_successful_generation_history_summary(
|
|
project.id,
|
|
path=db_path,
|
|
)
|
|
self.assertEqual(project.id, success_summary.project_id)
|
|
self.assertEqual(2, success_summary.successful_round_count)
|
|
self.assertEqual(3, success_summary.successful_image_count)
|
|
self.assertTrue(success_summary.latest_succeeded_at)
|
|
other_summary = image_studio.get_successful_generation_history_summary(
|
|
other_project.id,
|
|
path=db_path,
|
|
)
|
|
self.assertEqual(0, other_summary.successful_round_count)
|
|
self.assertEqual(0, other_summary.successful_image_count)
|
|
|
|
self.assert_removed(temp_dir)
|
|
|
|
def test_global_generation_rounds_filter_paginate_and_exclude_deleted_projects(self):
|
|
with self.make_temp_dir() as temp_dir:
|
|
db_path = os.path.join(temp_dir, "cmshopee.db")
|
|
db.init_db(db_path)
|
|
first_project = image_studio.create_or_get_project(
|
|
account_alias="main-shop",
|
|
account_name="主店",
|
|
account_slug="main-shop",
|
|
item_id="51100639510",
|
|
path=db_path,
|
|
)
|
|
second_project = image_studio.create_or_get_project(
|
|
account_alias="second-shop",
|
|
account_name="副店",
|
|
account_slug="second-shop",
|
|
item_id="51100639511",
|
|
path=db_path,
|
|
)
|
|
deleted_project = image_studio.create_or_get_project(
|
|
account_alias="deleted-shop",
|
|
account_slug="deleted-shop",
|
|
item_id="51100639512",
|
|
path=db_path,
|
|
)
|
|
first_source = image_studio.add_asset(
|
|
first_project.id,
|
|
image_studio.ASSET_KIND_ORIGINAL,
|
|
path=db_path,
|
|
)
|
|
second_source = image_studio.add_asset(
|
|
second_project.id,
|
|
image_studio.ASSET_KIND_ORIGINAL,
|
|
path=db_path,
|
|
)
|
|
deleted_source = image_studio.add_asset(
|
|
deleted_project.id,
|
|
image_studio.ASSET_KIND_ORIGINAL,
|
|
path=db_path,
|
|
)
|
|
first = image_studio.create_job(
|
|
first_project.id,
|
|
source_asset_id=first_source.id,
|
|
task_key="global-first-slot",
|
|
generation_round_key="first-round",
|
|
generation_slot_index=0,
|
|
path=db_path,
|
|
)
|
|
image_studio.update_job_status(first.id, "failed", path=db_path)
|
|
retry = image_studio.create_job(
|
|
first_project.id,
|
|
source_asset_id=first_source.id,
|
|
task_key="global-first-retry",
|
|
generation_round_key="first-round",
|
|
generation_slot_index=0,
|
|
path=db_path,
|
|
)
|
|
image_studio.update_job_status(retry.id, "succeeded", path=db_path)
|
|
second = image_studio.create_job(
|
|
second_project.id,
|
|
source_asset_id=second_source.id,
|
|
task_key="global-second-round",
|
|
generation_round_key="second-round",
|
|
generation_slot_index=0,
|
|
path=db_path,
|
|
)
|
|
image_studio.update_job_status(second.id, "succeeded", path=db_path)
|
|
deleted = image_studio.create_job(
|
|
deleted_project.id,
|
|
source_asset_id=deleted_source.id,
|
|
task_key="global-deleted-round",
|
|
generation_round_key="deleted-round",
|
|
generation_slot_index=0,
|
|
path=db_path,
|
|
)
|
|
image_studio.update_job_status(deleted.id, "succeeded", path=db_path)
|
|
legacy = image_studio.create_job(
|
|
first_project.id,
|
|
source_asset_id=first_source.id,
|
|
task_key="global-legacy-history",
|
|
path=db_path,
|
|
)
|
|
image_studio.update_job_status(legacy.id, "cancelled", path=db_path)
|
|
image_studio.set_current_generation_round(
|
|
first_project.id,
|
|
"first-round",
|
|
path=db_path,
|
|
)
|
|
image_studio.soft_delete_project(
|
|
deleted_project.id,
|
|
reason="测试软删除",
|
|
path=db_path,
|
|
)
|
|
|
|
rounds = image_studio.list_global_generation_rounds(path=db_path)
|
|
self.assertEqual(
|
|
{first_project.id, second_project.id},
|
|
{round_.project_id for round_ in rounds},
|
|
)
|
|
self.assertIsNone(rounds[0].generation_round_key)
|
|
self.assertTrue(rounds[0].is_legacy)
|
|
self.assertEqual(first_project.id, rounds[0].project_id)
|
|
self.assertEqual(second_project.id, rounds[1].project_id)
|
|
first_round = next(
|
|
round_
|
|
for round_ in rounds
|
|
if (
|
|
round_.project_id == first_project.id
|
|
and round_.generation_round_key == "first-round"
|
|
)
|
|
)
|
|
self.assertEqual("first-round", first_round.generation_round_key)
|
|
self.assertEqual("主店", first_round.account_name)
|
|
self.assertEqual(1, first_round.job_count)
|
|
self.assertEqual(1, first_round.slot_count)
|
|
self.assertEqual(1, first_round.retry_count)
|
|
self.assertEqual(1, first_round.succeeded_count)
|
|
self.assertEqual(0, first_round.failed_count)
|
|
self.assertTrue(first_round.is_current)
|
|
|
|
self.assertEqual(
|
|
[first_project.id, first_project.id],
|
|
[
|
|
round_.project_id
|
|
for round_ in image_studio.list_global_generation_rounds(
|
|
account_query="主店",
|
|
path=db_path,
|
|
)
|
|
],
|
|
)
|
|
self.assertEqual(
|
|
["main-shop", "second-shop"],
|
|
[
|
|
account.account_alias
|
|
for account in image_studio.list_global_history_accounts(path=db_path)
|
|
],
|
|
)
|
|
self.assertEqual(
|
|
[first_project.id, first_project.id],
|
|
[
|
|
round_.project_id
|
|
for round_ in image_studio.list_global_generation_rounds(
|
|
account_alias="main-shop",
|
|
path=db_path,
|
|
)
|
|
],
|
|
)
|
|
self.assertEqual(
|
|
[],
|
|
image_studio.list_global_generation_rounds(
|
|
account_alias="main",
|
|
path=db_path,
|
|
),
|
|
)
|
|
self.assertEqual(
|
|
[second_project.id],
|
|
[
|
|
round_.project_id
|
|
for round_ in image_studio.list_global_generation_rounds(
|
|
item_query="51100639511",
|
|
path=db_path,
|
|
)
|
|
],
|
|
)
|
|
self.assertEqual(
|
|
[first_project.id, first_project.id],
|
|
[
|
|
round_.project_id
|
|
for round_ in image_studio.list_global_generation_rounds(
|
|
project_id=first_project.id,
|
|
path=db_path,
|
|
)
|
|
],
|
|
)
|
|
self.assertEqual(1, len(image_studio.list_global_generation_rounds(limit=1, path=db_path)))
|
|
self.assertEqual(1, len(image_studio.list_global_generation_rounds(limit=1, offset=1, path=db_path)))
|
|
|
|
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()
|