feat(product-suite): persist current generation rounds
This commit is contained in:
chengma
+171
-1
@@ -62,6 +62,7 @@ class ImageStudioTests(TempDirMixin, unittest.TestCase):
|
||||
"target_main_count",
|
||||
"target_detail_count",
|
||||
"suite_settings_json",
|
||||
"current_generation_round_key",
|
||||
"deleted_at",
|
||||
}.issubset(projects_columns)
|
||||
)
|
||||
@@ -69,7 +70,20 @@ class ImageStudioTests(TempDirMixin, unittest.TestCase):
|
||||
row["name"]
|
||||
for row in conn.execute("PRAGMA table_info(image_studio_jobs)").fetchall()
|
||||
}
|
||||
self.assertIn("recovery_action", jobs_columns)
|
||||
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()
|
||||
|
||||
@@ -820,6 +834,162 @@ class ImageStudioTests(TempDirMixin, unittest.TestCase):
|
||||
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):
|
||||
|
||||
@@ -1358,6 +1358,220 @@ class ProductSuiteGuiTests(TempDirMixin, unittest.TestCase):
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
def test_generation_round_restores_after_project_rebind_and_retry_keeps_slot(self):
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
config = self._config(temp_dir)
|
||||
project, sources = self._create_project_with_assets(temp_dir, config, 1)
|
||||
source = sources[0]
|
||||
round_key = "persisted-round"
|
||||
first = image_studio.create_job(
|
||||
project.id,
|
||||
source_asset_id=source.id,
|
||||
job_type="白底图",
|
||||
prompt="第一张",
|
||||
generation_round_key=round_key,
|
||||
generation_slot_index=0,
|
||||
path=config["db_path"],
|
||||
)
|
||||
first = image_studio.update_job_status(
|
||||
first.id,
|
||||
"succeeded",
|
||||
path=config["db_path"],
|
||||
)
|
||||
failed = image_studio.create_job(
|
||||
project.id,
|
||||
source_asset_id=source.id,
|
||||
job_type="场景图",
|
||||
prompt="第二张",
|
||||
generation_round_key=round_key,
|
||||
generation_slot_index=1,
|
||||
path=config["db_path"],
|
||||
)
|
||||
failed = image_studio.update_job_status(
|
||||
failed.id,
|
||||
"failed",
|
||||
path=config["db_path"],
|
||||
)
|
||||
image_studio.set_current_generation_round(
|
||||
project.id,
|
||||
round_key,
|
||||
path=config["db_path"],
|
||||
)
|
||||
|
||||
tab = ProductSuiteTab(config=config, db_path=config["db_path"])
|
||||
self.addCleanup(tab.close)
|
||||
state = tab._displayed_state
|
||||
state.account_alias = "alias-a"
|
||||
state.item_id = project.item_id
|
||||
tab._bind_project(state, load_existing=True)
|
||||
self.assertEqual(round_key, state.current_generation_round_key)
|
||||
self.assertEqual([first.id, failed.id], state.current_job_ids)
|
||||
|
||||
retry = image_studio.create_job(
|
||||
project.id,
|
||||
source_asset_id=source.id,
|
||||
job_type=failed.job_type,
|
||||
prompt=failed.prompt,
|
||||
generation_round_key=failed.generation_round_key,
|
||||
generation_slot_index=failed.generation_slot_index,
|
||||
path=config["db_path"],
|
||||
)
|
||||
image_studio.update_job_status(retry.id, "succeeded", path=config["db_path"])
|
||||
|
||||
reopened = ProductSuiteTab(config=config, db_path=config["db_path"])
|
||||
self.addCleanup(reopened.close)
|
||||
reopened_state = reopened._displayed_state
|
||||
reopened_state.account_alias = "alias-a"
|
||||
reopened_state.item_id = project.item_id
|
||||
reopened._bind_project(reopened_state, load_existing=True)
|
||||
self.assertEqual(round_key, reopened_state.current_generation_round_key)
|
||||
self.assertEqual([first.id, retry.id], reopened_state.current_job_ids)
|
||||
self.assertEqual(
|
||||
[first.id, retry.id],
|
||||
[job.id for job in reopened._jobs_for_state(reopened_state)],
|
||||
)
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
def test_new_generation_round_promotes_partial_success_and_keeps_previous_on_failure(self):
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
config = self._config(temp_dir)
|
||||
project, sources = self._create_project_with_assets(temp_dir, config, 1)
|
||||
source = sources[0]
|
||||
old_round = "old-current-round"
|
||||
old_job = image_studio.create_job(
|
||||
project.id,
|
||||
source_asset_id=source.id,
|
||||
job_type="白底图",
|
||||
prompt="旧结果",
|
||||
generation_round_key=old_round,
|
||||
generation_slot_index=0,
|
||||
path=config["db_path"],
|
||||
)
|
||||
old_job = image_studio.update_job_status(
|
||||
old_job.id,
|
||||
"succeeded",
|
||||
path=config["db_path"],
|
||||
)
|
||||
image_studio.set_current_generation_round(
|
||||
project.id,
|
||||
old_round,
|
||||
path=config["db_path"],
|
||||
)
|
||||
|
||||
tab = ProductSuiteTab(config=config, db_path=config["db_path"])
|
||||
self.addCleanup(tab.close)
|
||||
state = tab._displayed_state
|
||||
state.account_alias = "alias-a"
|
||||
state.item_id = project.item_id
|
||||
tab._bind_project(state, load_existing=True)
|
||||
messages = []
|
||||
tab._message = lambda title, message, **kwargs: messages.append(title)
|
||||
|
||||
first_run_jobs = []
|
||||
|
||||
def partial_success(jobs, **kwargs):
|
||||
first_run_jobs[:] = list(jobs)
|
||||
image_studio.update_job_status(
|
||||
first_run_jobs[0].id,
|
||||
"succeeded",
|
||||
path=config["db_path"],
|
||||
)
|
||||
image_studio.update_job_status(
|
||||
first_run_jobs[1].id,
|
||||
"failed",
|
||||
path=config["db_path"],
|
||||
)
|
||||
return {"total": 2, "success": 1, "failed": 1, "cancelled": 0}
|
||||
|
||||
specs = [
|
||||
{"source_asset_id": source.id, "job_type": "白底图", "prompt": "新图1"},
|
||||
{"source_asset_id": source.id, "job_type": "场景图", "prompt": "新图2"},
|
||||
]
|
||||
with mock.patch(
|
||||
"app.gui.workers.image_studio_generation.run_jobs",
|
||||
side_effect=partial_success,
|
||||
):
|
||||
self.assertTrue(tab.start_generation(state, specs=specs))
|
||||
first_thread = state.thread
|
||||
deadline = time.monotonic() + 3
|
||||
while state.worker is not None and time.monotonic() < deadline:
|
||||
QTest.qWait(20)
|
||||
self.app.processEvents()
|
||||
if first_thread is not None:
|
||||
try:
|
||||
deadline = time.monotonic() + 3
|
||||
while first_thread.isRunning() and time.monotonic() < deadline:
|
||||
QTest.qWait(20)
|
||||
self.app.processEvents()
|
||||
self.assertFalse(first_thread.isRunning())
|
||||
except RuntimeError:
|
||||
pass
|
||||
|
||||
new_round = image_studio.get_current_generation_round(
|
||||
project.id,
|
||||
path=config["db_path"],
|
||||
)
|
||||
self.assertNotEqual(old_round, new_round)
|
||||
self.assertEqual(
|
||||
[(new_round, 0), (new_round, 1)],
|
||||
[
|
||||
(job.generation_round_key, job.generation_slot_index)
|
||||
for job in first_run_jobs
|
||||
],
|
||||
)
|
||||
self.assertEqual(
|
||||
[job.id for job in first_run_jobs],
|
||||
state.current_job_ids,
|
||||
)
|
||||
|
||||
def all_failed(jobs, **kwargs):
|
||||
for job in jobs:
|
||||
image_studio.update_job_status(
|
||||
job.id,
|
||||
"failed",
|
||||
path=config["db_path"],
|
||||
)
|
||||
return {"total": 2, "success": 0, "failed": 2, "cancelled": 0}
|
||||
|
||||
with mock.patch(
|
||||
"app.gui.workers.image_studio_generation.run_jobs",
|
||||
side_effect=all_failed,
|
||||
):
|
||||
self.assertTrue(tab.start_generation(state, specs=specs))
|
||||
second_thread = state.thread
|
||||
deadline = time.monotonic() + 3
|
||||
while state.worker is not None and time.monotonic() < deadline:
|
||||
QTest.qWait(20)
|
||||
self.app.processEvents()
|
||||
if second_thread is not None:
|
||||
try:
|
||||
deadline = time.monotonic() + 3
|
||||
while second_thread.isRunning() and time.monotonic() < deadline:
|
||||
QTest.qWait(20)
|
||||
self.app.processEvents()
|
||||
self.assertFalse(second_thread.isRunning())
|
||||
except RuntimeError:
|
||||
pass
|
||||
|
||||
self.assertEqual(
|
||||
new_round,
|
||||
image_studio.get_current_generation_round(
|
||||
project.id,
|
||||
path=config["db_path"],
|
||||
),
|
||||
)
|
||||
self.assertEqual(
|
||||
[job.id for job in first_run_jobs],
|
||||
state.current_job_ids,
|
||||
)
|
||||
self.assertEqual(
|
||||
["商品套图生成完成", "商品套图生成完成"],
|
||||
messages,
|
||||
)
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
def test_retry_tracks_only_new_job_and_preserves_history_view(self):
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
config = self._config(temp_dir)
|
||||
|
||||
+56
-1
@@ -1,4 +1,5 @@
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest import mock
|
||||
@@ -7,7 +8,7 @@ os.environ.setdefault("QT_QPA_PLATFORM", "offscreen")
|
||||
|
||||
from _helpers import REPO_ROOT # noqa: F401
|
||||
|
||||
from app import image_studio_images, workers
|
||||
from app import db, image_studio, image_studio_images, workers
|
||||
|
||||
if workers.QT_IMPORT_ERROR is not None:
|
||||
raise unittest.SkipTest("PySide6 未安装")
|
||||
@@ -19,6 +20,7 @@ from app.workers import BaseWorker, run_worker
|
||||
from app.gui.workers import (
|
||||
ImageStudioDownloadOriginalWorker,
|
||||
ImageStudioPullImagesWorker,
|
||||
ProductSuiteGenerateWorker,
|
||||
)
|
||||
|
||||
|
||||
@@ -196,6 +198,59 @@ class WorkerTests(unittest.TestCase):
|
||||
|
||||
self.assertEqual({"asset_id": 12, "cancelled": True}, summary)
|
||||
|
||||
def test_product_suite_worker_writes_generation_round_and_stable_slots(self):
|
||||
with tempfile.TemporaryDirectory() 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,
|
||||
)
|
||||
source = image_studio.add_asset(
|
||||
project.id,
|
||||
image_studio.ASSET_KIND_ORIGINAL,
|
||||
path=db_path,
|
||||
)
|
||||
worker = ProductSuiteGenerateWorker(
|
||||
project.id,
|
||||
[
|
||||
{
|
||||
"source_asset_id": source.id,
|
||||
"job_type": "白底图",
|
||||
"prompt": "第一张",
|
||||
"generation_round_key": "worker-round",
|
||||
"generation_slot_index": 0,
|
||||
},
|
||||
{
|
||||
"source_asset_id": source.id,
|
||||
"job_type": "场景图",
|
||||
"prompt": "第二张",
|
||||
"generation_round_key": "worker-round",
|
||||
"generation_slot_index": 1,
|
||||
},
|
||||
],
|
||||
generation_round_key="worker-round",
|
||||
db_path=db_path,
|
||||
)
|
||||
|
||||
with mock.patch(
|
||||
"app.gui.workers.image_studio_generation.run_jobs",
|
||||
return_value={"total": 2, "success": 0, "failed": 0, "cancelled": 0},
|
||||
):
|
||||
summary = worker.execute()
|
||||
|
||||
jobs = [image_studio.get_job(job_id, path=db_path) for job_id in worker.job_ids]
|
||||
self.assertEqual("worker-round", summary["generation_round_key"])
|
||||
self.assertEqual(
|
||||
[("worker-round", 0), ("worker-round", 1)],
|
||||
[
|
||||
(job.generation_round_key, job.generation_slot_index)
|
||||
for job in jobs
|
||||
],
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
Reference in New Issue
Block a user