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