feat(product-suite): persist current generation rounds
This commit is contained in:
chengma
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user