332 lines
13 KiB
Python
332 lines
13 KiB
Python
import os
|
|
import sys
|
|
import unittest
|
|
from types import SimpleNamespace
|
|
|
|
sys.path.insert(0, os.path.dirname(__file__))
|
|
|
|
from _helpers import TempDirMixin
|
|
|
|
from app import db, image_studio
|
|
|
|
|
|
class ImageStudioTests(TempDirMixin, unittest.TestCase):
|
|
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",
|
|
"target_main_count",
|
|
"target_detail_count",
|
|
"deleted_at",
|
|
}.issubset(projects_columns)
|
|
)
|
|
finally:
|
|
conn.close()
|
|
|
|
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)
|
|
|
|
updated = image_studio.update_project_prompt(project.id, "二次提示词", path=db_path)
|
|
self.assertEqual("二次提示词", updated.draft_prompt)
|
|
|
|
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_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_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)
|
|
|
|
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([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.assertIsNotNone(succeeded.finished_at)
|
|
self.assertEqual([], image_studio.list_resumable_jobs(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_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)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|