feat: snapshot product suite reference assets

This commit is contained in:
chengma
2026-07-17 17:03:40 +08:00
parent d11dc752b5
commit 1fabd8afcf
12 changed files with 252 additions and 5 deletions
+43
View File
@@ -75,6 +75,7 @@ class ImageStudioTests(TempDirMixin, unittest.TestCase):
"recovery_action",
"generation_round_key",
"generation_slot_index",
"reference_asset_ids",
}.issubset(jobs_columns)
)
indexes = {
@@ -843,6 +844,48 @@ class ImageStudioTests(TempDirMixin, unittest.TestCase):
self.assertTrue(
all(job.generation_round_key is None for job in legacy_jobs)
)
self.assertTrue(
all(job.reference_asset_ids is None for job in legacy_jobs)
)
self.assert_removed(temp_dir)
def test_create_job_freezes_valid_reference_asset_ids(self):
with self.make_temp_dir() as temp_dir:
config = self._config(temp_dir)
db.init_db(config["db_path"])
project = image_studio.create_or_get_project(
account_alias="店铺",
account_slug="shop",
item_id="51100639510",
path=config["db_path"],
)
source = image_studio.add_asset(
project.id,
image_studio.ASSET_KIND_ORIGINAL,
path=config["db_path"],
)
reference = image_studio.add_asset(
project.id,
image_studio.ASSET_KIND_ORIGINAL,
path=config["db_path"],
)
job = image_studio.create_job(
project.id,
source_asset_id=source.id,
reference_asset_ids=[reference.id],
path=config["db_path"],
)
self.assertEqual("[%d]" % reference.id, job.reference_asset_ids)
self.assertEqual([reference.id], image_studio.job_reference_asset_ids(job))
with self.assertRaisesRegex(db.DbError, "不能包含主图"):
image_studio.create_job(
project.id,
source_asset_id=source.id,
reference_asset_ids=[source.id],
path=config["db_path"],
)
self.assert_removed(temp_dir)
+83
View File
@@ -246,6 +246,89 @@ class ImageStudioGenerationTests(TempDirMixin, unittest.TestCase):
self.assert_removed(temp_dir)
def test_job_reference_snapshot_submits_ordered_images(self):
with self.make_temp_dir() as temp_dir:
cfg, project, source = self._project_source(temp_dir)
reference_path = os.path.join(temp_dir, "reference.png")
with open(reference_path, "wb") as fh:
fh.write(self._png_bytes())
reference = image_studio.add_asset(
project.id,
image_studio.ASSET_KIND_ORIGINAL,
local_path=reference_path,
source_order=2,
path=cfg["db_path"],
)
job = image_studio.create_job(
project.id,
source_asset_id=source.id,
reference_asset_ids=[reference.id],
prompt="多图提示词",
path=cfg["db_path"],
)
submitted = []
def fake_submit(method, url, api_key, **kwargs):
submitted.append(kwargs["payload"])
return {"task_id": "multi-image-task", "status": "queued"}
with mock.patch("app.image_studio_generation._runtime", return_value=self._runtime()), \
mock.patch("app.image_studio_generation.ai._cmhub_call_with_retry", side_effect=fake_submit), \
mock.patch(
"app.image_studio_generation.ai._cmhub_call_once",
return_value={
"task_id": "multi-image-task",
"status": "succeeded",
"result": {"image_url": "https://cdn.example.com/multi.png"},
},
), \
mock.patch(
"app.image_studio_generation.ai._download_cmhub_image_with_retry",
return_value=(self._png_bytes(), 0.1),
):
summary = image_studio_generation.run_jobs(
[job],
config=cfg,
path=cfg["db_path"],
)
self.assertEqual(1, summary["success"])
self.assertEqual(2, len(submitted[0]["images"]))
self.assertNotIn("image_base64", submitted[0])
self.assert_removed(temp_dir)
def test_missing_reference_snapshot_fails_without_submitting(self):
with self.make_temp_dir() as temp_dir:
cfg, project, source = self._project_source(temp_dir)
reference = image_studio.add_asset(
project.id,
image_studio.ASSET_KIND_ORIGINAL,
local_path=os.path.join(temp_dir, "missing-reference.png"),
source_order=2,
path=cfg["db_path"],
)
job = image_studio.create_job(
project.id,
source_asset_id=source.id,
reference_asset_ids=[reference.id],
prompt="多图提示词",
path=cfg["db_path"],
)
with mock.patch("app.image_studio_generation._runtime", return_value=self._runtime()), \
mock.patch("app.image_studio_generation.ai._cmhub_call_with_retry") as submit:
summary = image_studio_generation.run_jobs(
[job],
config=cfg,
path=cfg["db_path"],
)
self.assertEqual(1, summary["failed"])
self.assertIn("参考图尚未下载", summary["jobs"][0]["error"])
submit.assert_not_called()
self.assert_removed(temp_dir)
def test_stop_after_download_discards_temporary_result(self):
with self.make_temp_dir() as temp_dir:
cfg, project, source = self._project_source(temp_dir)
+21
View File
@@ -99,6 +99,7 @@ class ProductSuiteTests(unittest.TestCase):
self.assertEqual(3, len(specs))
self.assertEqual([11, 11, 12], [spec["source_asset_id"] for spec in specs])
self.assertEqual([[], [], []], [spec["reference_asset_ids"] for spec in specs])
self.assertEqual(["白底图", "场景图", "场景图"], [spec["job_type"] for spec in specs])
for spec in specs:
self.assertIn("平台:Shopee", spec["prompt"])
@@ -135,6 +136,26 @@ class ProductSuiteTests(unittest.TestCase):
self.assertLess(white_prompt.index("尺码信息规则"), white_prompt.index(reference_rule))
self.assertLess(white_prompt.index(reference_rule), white_prompt.index("商品卖点与要求"))
def test_job_specs_freeze_first_image_references_when_not_per_image_primary(self):
settings = product_suite.default_suite_settings()
settings.update(
{
"per_image_primary": False,
"categories": {"白底图": 1, "场景图": 1, "卖点图": 0},
}
)
specs = product_suite.build_job_specs(
[SimpleNamespace(id=11), SimpleNamespace(id=12), SimpleNamespace(id=13)],
"卖点",
settings,
"51100639510",
template_text=prompts.load_default_product_suite_prompt(),
)
self.assertEqual(2, len(specs))
self.assertEqual([11, 11], [spec["source_asset_id"] for spec in specs])
self.assertEqual([[12, 13], [12, 13]], [spec["reference_asset_ids"] for spec in specs])
def test_reference_rule_follows_per_image_primary_setting(self):
template = prompts.load_default_product_suite_prompt()
settings = product_suite.default_suite_settings()