chengma/cmshoppe
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(ai-studio): separate image pool from task states

Browse Source
This commit is contained in:
chengma
2026-07-13 09:54:08 +08:00
parent 42fe233ed9
commit b1c050f633
9 changed files with 500 additions and 59 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_gui.py
+127 -3
View File
@@ -928,19 +928,143 @@ class GuiTests(TempDirMixin, unittest.TestCase):
tab._select_project(project.id)
self.assertEqual("继续查询任务", tab.resume_button.text())
self.assertEqual(1, tab.pool_grid.count())
self.assertTrue(
all(
tab.pool_grid.item(row).data(gui.Qt.UserRole)["type"] == "asset"
for row in range(tab.pool_grid.count())
)
)
self.assertFalse(tab.job_section.isHidden())
statuses = [
tab.pool_grid.item(row).text()
for row in range(tab.pool_grid.count())
if tab.pool_grid.item(row).data(gui.Qt.UserRole)["type"] == "job"
tab.job_grid.item(row).text()
for row in range(tab.job_grid.count())
]
self.assertEqual(1, len(statuses))
self.assertIn("已提交", statuses[0])
self.assertIn("可继续查询", statuses[0])
self.assertIn("扣点2", statuses[0])
self.assertIn("余额88", statuses[0])
self.assertIn("call_id=call-1", statuses[0])
self.assert_removed(temp_dir)
def test_image_studio_tab_separates_usable_assets_from_job_states(self):
with self.make_temp_dir() as temp_dir:
cfg = self.make_config(temp_dir)
db.init_db(cfg["db_path"])
project = image_studio.create_or_get_project(
account_alias="alias-a",
account_slug="alias_a",
item_id="51100639510",
path=cfg["db_path"],
)
other_project = image_studio.create_or_get_project(
account_alias="alias-b",
account_slug="alias_b",
item_id="51100639511",
path=cfg["db_path"],
)
source = image_studio.add_asset(
project.id,
"original",
local_path=self.write_test_image(os.path.join(temp_dir, "source.jpg")),
path=cfg["db_path"],
)
resumable = image_studio.create_job(
project.id,
source_asset_id=source.id,
task_key="resume-job",
path=cfg["db_path"],
)
image_studio.set_job_submitted(resumable.id, "remote-resume", path=cfg["db_path"])
image_studio.update_job_status(
resumable.id,
"failed",
error="下载失败：https://secret.example.com/api/v1/image?token=private",
path=cfg["db_path"],
)
terminal = image_studio.create_job(
project.id,
source_asset_id=source.id,
task_key="terminal-job",
path=cfg["db_path"],
)
image_studio.set_job_submitted(terminal.id, "remote-terminal", path=cfg["db_path"])
image_studio.update_job_status(
terminal.id,
"expired",
error="上游任务已过期",
recovery_action=image_studio.JOB_RECOVERY_REGENERATE,
path=cfg["db_path"],
)
stopped = image_studio.create_job(
project.id,
source_asset_id=source.id,
task_key="stopped-job",
path=cfg["db_path"],
)
image_studio.update_job_status(
stopped.id,
"cancelled",
error="用户停止",
recovery_action=image_studio.JOB_RECOVERY_REGENERATE,
path=cfg["db_path"],
)
tab = ImageStudioTab(config=cfg, db_path=cfg["db_path"])
self.addCleanup(tab.close)
tab._select_project(project.id)
self.assertEqual(1, tab.pool_grid.count())
self.assertTrue(
all(
tab.pool_grid.item(row).data(gui.Qt.UserRole)["type"] == "asset"
for row in range(tab.pool_grid.count())
)
)
self.assertFalse(tab.job_section.isHidden())
self.assertEqual(3, tab.job_grid.count())
job_items = [tab.job_grid.item(row) for row in range(tab.job_grid.count())]
self.assertTrue(
all(item.data(gui.Qt.UserRole)["draggable"] is False for item in job_items)
)
details = "\n".join(item.text() for item in job_items)
tooltips = "\n".join(item.toolTip() for item in job_items)
self.assertIn("生成失败", details)
self.assertIn("任务过期", details)
self.assertIn("已停止", details)
self.assertIn("可继续查询", details)
self.assertIn("需要重新生成，可能再次扣点", details)
self.assertNotIn("https://", details)
self.assertNotIn("secret.example", tooltips)
failed_item = next(item for item in job_items if "生成失败" in item.text())
self.assertEqual("#ffebe9", failed_item.background().color().name())
tab._select_project(other_project.id)
self.assertEqual(0, tab.job_grid.count())
self.assertTrue(tab.job_section.isHidden())
tab._select_project(project.id)
output = image_studio.add_asset(
project.id,
"generated_main",
parent_asset_id=source.id,
local_path=self.write_test_image(os.path.join(temp_dir, "generated.jpg")),
path=cfg["db_path"],
)
image_studio.update_job_status(
resumable.id,
"succeeded",
output_asset_id=output.id,
path=cfg["db_path"],
)
tab.refresh_project_assets()
self.assertEqual(2, tab.pool_grid.count())
self.assertEqual(2, tab.job_grid.count())
self.assert_removed(temp_dir)
def test_image_studio_generation_log_uses_cmhub_tier_summary(self):
with self.make_temp_dir() as temp_dir:
cfg = self.make_config(temp_dir)
tests/test_image_studio.py
+95
View File
@@ -62,6 +62,11 @@ class ImageStudioTests(TempDirMixin, unittest.TestCase):
"deleted_at",
}.issubset(projects_columns)
)
jobs_columns = {
row["name"]
for row in conn.execute("PRAGMA table_info(image_studio_jobs)").fetchall()
}
self.assertIn("recovery_action", jobs_columns)
finally:
conn.close()
@@ -305,6 +310,7 @@ class ImageStudioTests(TempDirMixin, unittest.TestCase):
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,
@@ -317,6 +323,7 @@ class ImageStudioTests(TempDirMixin, unittest.TestCase):
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(
@@ -346,9 +353,30 @@ class ImageStudioTests(TempDirMixin, unittest.TestCase):
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),
)
with self.assertRaises(db.DbError):
image_studio.create_job(
project.id,
@@ -358,6 +386,73 @@ class ImageStudioTests(TempDirMixin, unittest.TestCase):
self.assert_removed(temp_dir)
def test_init_db_migrates_legacy_image_studio_job_recovery_action(self):
with self.make_temp_dir() as temp_dir:
db_path = os.path.join(temp_dir, "legacy.db")
conn = db.connect(db_path)
try:
conn.execute(
"""
CREATE TABLE image_studio_jobs (
id INTEGER PRIMARY KEY,
project_id INTEGER NOT NULL,
source_asset_id INTEGER,
output_asset_id INTEGER,
generation_source TEXT NOT NULL DEFAULT 'cmhub',
provider TEXT NOT NULL DEFAULT 'cmhub',
job_type TEXT NOT NULL,
task_key TEXT NOT NULL UNIQUE,
task_id TEXT,
status TEXT NOT NULL DEFAULT 'pending',
prompt TEXT,
error TEXT,
attempts INTEGER NOT NULL DEFAULT 0,
points_cost INTEGER,
points_balance INTEGER,
call_id TEXT,
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL,
submitted_at TEXT,
finished_at TEXT
)
"""
)
conn.executemany(
"""
INSERT INTO image_studio_jobs
(id, project_id, job_type, task_key, task_id, status, created_at, updated_at)
VALUES (?, 1, 'main', ?, ?, ?, '2026-07-13T00:00:00', '2026-07-13T00:00:00')
""",
[
(1, "legacy-submitted", "task-submitted", "submitted"),
(2, "legacy-running", "task-running", "running"),
(3, "legacy-failed", "task-failed", "failed"),
(4, "legacy-success", "task-success", "succeeded"),
],
)
conn.commit()
finally:
conn.close()
db.init_db(db_path)
conn = db.connect(db_path)
try:
recovery_actions = {
row["task_key"]: row["recovery_action"]
for row in conn.execute(
"SELECT task_key, recovery_action FROM image_studio_jobs ORDER BY id"
).fetchall()
}
finally:
conn.close()
self.assertEqual(image_studio.JOB_RECOVERY_RESUME, recovery_actions["legacy-submitted"])
self.assertEqual(image_studio.JOB_RECOVERY_RESUME, recovery_actions["legacy-running"])
self.assertEqual(image_studio.JOB_RECOVERY_REGENERATE, recovery_actions["legacy-failed"])
self.assertEqual(image_studio.JOB_RECOVERY_NONE, recovery_actions["legacy-success"])
self.assert_removed(temp_dir)
def test_selections_are_consecutive_unique_and_replaceable(self):
with self.make_temp_dir() as temp_dir:
db_path = os.path.join(temp_dir, "cmshopee.db")
tests/test_image_studio_generation.py
+14
View File
@@ -256,6 +256,9 @@ class ImageStudioGenerationTests(TempDirMixin, unittest.TestCase):
path=cfg["db_path"],
)
failed_job = image_studio.get_job(job.id, path=cfg["db_path"])
self.assertEqual(image_studio.JOB_RECOVERY_RESUME, failed_job.recovery_action)
resumable = image_studio.list_resumable_jobs(
path=cfg["db_path"],
project_id=project.id,
@@ -291,6 +294,7 @@ class ImageStudioGenerationTests(TempDirMixin, unittest.TestCase):
self.assertEqual("succeeded", updated.status)
self.assertEqual("cmhub-task-download", updated.task_id)
self.assertEqual("call-download", updated.call_id)
self.assertEqual(image_studio.JOB_RECOVERY_NONE, updated.recovery_action)
assets = image_studio.list_assets(project.id, kind="generated_main", path=cfg["db_path"])
self.assertEqual(1, len(assets))
@@ -339,6 +343,16 @@ class ImageStudioGenerationTests(TempDirMixin, unittest.TestCase):
self.assertEqual(1, summary["failed"])
assets = image_studio.list_assets(project.id, kind="generated_main", path=cfg["db_path"])
self.assertEqual(1, len(assets))
failed_jobs = [
job
for job in image_studio.list_resumable_jobs(
path=cfg["db_path"],
project_id=project.id,
include_failed_downloads=True,
)
if job.status == "failed"
]
self.assertEqual([], failed_jobs)
self.assert_removed(temp_dir)