feat: complete T-533 incremental generation
This commit is contained in:
+129
-14
@@ -757,35 +757,150 @@ class AITests(TempDirMixin, unittest.TestCase):
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
def test_generate_batch_retries_failed_generation_record_after_existing_result(self):
|
||||
def test_generate_batch_fills_missing_cover_without_regenerating_title(self):
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
cfg = self._config()
|
||||
cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
|
||||
cfg["image_dir"] = os.path.join(temp_dir, "images")
|
||||
cfg["ai"]["generate_cover"] = True
|
||||
batch_id, tasks = self._collected_tasks(temp_dir, cfg, ["旧标题"])
|
||||
db.set_generated(tasks[0].id, "手动标题", None, path=cfg["db_path"])
|
||||
cover_only_task = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])[0]
|
||||
progress = []
|
||||
|
||||
def fake_cover(cover_prompt, old_cover_path, out_path, **kwargs):
|
||||
os.makedirs(os.path.dirname(out_path), exist_ok=True)
|
||||
with open(out_path, "wb") as fh:
|
||||
fh.write(b"jpeg")
|
||||
return out_path
|
||||
|
||||
with mock.patch("app.ai.gen_title") as gen_title, \
|
||||
mock.patch("app.ai.gen_cover", side_effect=fake_cover) as gen_cover:
|
||||
summary = ai.generate_batch(
|
||||
[cover_only_task],
|
||||
{"title": "标题提示", "cover": "封面 {新标题}"},
|
||||
ai_cfg={"config": cfg, "db_path": cfg["db_path"]},
|
||||
on_progress=progress.append,
|
||||
)
|
||||
|
||||
self.assertTrue(summary["ok"])
|
||||
self.assertEqual(1, summary["total"])
|
||||
self.assertEqual(0, summary["title_total"])
|
||||
self.assertEqual(0, summary["title_done"])
|
||||
self.assertEqual(1, summary["cover_total"])
|
||||
self.assertEqual(1, summary["cover_done"])
|
||||
self.assertEqual(1, summary["generated_done"])
|
||||
self.assertEqual(0, progress[-1]["title_total"])
|
||||
self.assertEqual(1, progress[-1]["cover_total"])
|
||||
gen_title.assert_not_called()
|
||||
gen_cover.assert_called_once()
|
||||
updated = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])[0]
|
||||
self.assertEqual("generated", updated.stage)
|
||||
self.assertEqual("success", updated.status)
|
||||
self.assertEqual("手动标题", updated.new_title)
|
||||
self.assertTrue(os.path.exists(updated.new_cover_path))
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
def test_generate_batch_handles_mixed_title_and_cover_gaps(self):
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
cfg = self._config()
|
||||
cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
|
||||
cfg["image_dir"] = os.path.join(temp_dir, "images")
|
||||
cfg["ai"]["generate_cover"] = True
|
||||
batch_id, tasks = self._collected_tasks(temp_dir, cfg, ["旧标题A", "旧标题B"])
|
||||
db.set_generated(tasks[1].id, "已有标题B", None, path=cfg["db_path"])
|
||||
mixed_tasks = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])
|
||||
cover_prompts = []
|
||||
|
||||
def fake_title(title_prompt, old_title, **kwargs):
|
||||
return "新" + old_title
|
||||
|
||||
def fake_cover(cover_prompt, old_cover_path, out_path, **kwargs):
|
||||
cover_prompts.append(cover_prompt)
|
||||
os.makedirs(os.path.dirname(out_path), exist_ok=True)
|
||||
with open(out_path, "wb") as fh:
|
||||
fh.write(b"jpeg")
|
||||
return out_path
|
||||
|
||||
with mock.patch("app.ai.gen_title", side_effect=fake_title) as gen_title, \
|
||||
mock.patch("app.ai.gen_cover", side_effect=fake_cover) as gen_cover:
|
||||
summary = ai.generate_batch(
|
||||
mixed_tasks,
|
||||
{"title": "标题提示", "cover": "封面 {新标题}"},
|
||||
ai_cfg={"config": cfg, "db_path": cfg["db_path"]},
|
||||
)
|
||||
|
||||
self.assertTrue(summary["ok"])
|
||||
self.assertEqual(2, summary["total"])
|
||||
self.assertEqual(1, summary["title_total"])
|
||||
self.assertEqual(1, summary["title_done"])
|
||||
self.assertEqual(2, summary["cover_total"])
|
||||
self.assertEqual(2, summary["cover_done"])
|
||||
self.assertEqual(2, summary["generated_done"])
|
||||
self.assertEqual(1, gen_title.call_count)
|
||||
self.assertEqual(2, gen_cover.call_count)
|
||||
updated = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])
|
||||
by_old_title = {task.old_title: task for task in updated}
|
||||
self.assertEqual("新旧标题A", by_old_title["旧标题A"].new_title)
|
||||
self.assertEqual("已有标题B", by_old_title["旧标题B"].new_title)
|
||||
self.assertTrue(all(os.path.exists(task.new_cover_path) for task in updated))
|
||||
self.assertIn("新旧标题A", "\n".join(cover_prompts))
|
||||
self.assertIn("已有标题B", "\n".join(cover_prompts))
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
def test_generate_batch_ignores_title_only_task_when_cover_disabled(self):
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
cfg = self._config()
|
||||
cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
|
||||
cfg["image_dir"] = os.path.join(temp_dir, "images")
|
||||
cfg["ai"]["generate_cover"] = False
|
||||
batch_id, tasks = self._collected_tasks(temp_dir, cfg, ["旧标题"])
|
||||
db.set_generated(tasks[0].id, "旧AI标题", "old-new.jpg", path=cfg["db_path"])
|
||||
db.mark_failed(tasks[0].id, "generate", "上次生成失败", path=cfg["db_path"])
|
||||
retry_task = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])[0]
|
||||
db.set_generated(tasks[0].id, "已有标题", None, path=cfg["db_path"])
|
||||
title_only_task = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])[0]
|
||||
|
||||
with mock.patch("app.ai.gen_title", return_value="重新生成标题") as gen_title, \
|
||||
with mock.patch("app.ai.gen_title") as gen_title, \
|
||||
mock.patch("app.ai.gen_cover") as gen_cover:
|
||||
summary = ai.generate_batch(
|
||||
[retry_task],
|
||||
[title_only_task],
|
||||
{"title": "标题提示", "cover": "封面"},
|
||||
ai_cfg={"config": cfg, "db_path": cfg["db_path"]},
|
||||
)
|
||||
|
||||
self.assertTrue(summary["ok"])
|
||||
self.assertEqual(1, summary["total"])
|
||||
self.assertEqual(1, summary["generated_done"])
|
||||
gen_title.assert_called_once()
|
||||
self.assertEqual(0, summary["total"])
|
||||
self.assertEqual(0, summary["title_total"])
|
||||
self.assertEqual(0, summary["cover_total"])
|
||||
gen_title.assert_not_called()
|
||||
gen_cover.assert_not_called()
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
def test_generate_batch_ignores_complete_generated_task(self):
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
cfg = self._config()
|
||||
cfg["db_path"] = os.path.join(temp_dir, "cmshopee.db")
|
||||
cfg["image_dir"] = os.path.join(temp_dir, "images")
|
||||
cfg["ai"]["generate_cover"] = True
|
||||
batch_id, tasks = self._collected_tasks(temp_dir, cfg, ["旧标题"])
|
||||
db.set_generated(tasks[0].id, "已有标题", "new.jpg", path=cfg["db_path"])
|
||||
complete_task = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])[0]
|
||||
|
||||
with mock.patch("app.ai.gen_title") as gen_title, \
|
||||
mock.patch("app.ai.gen_cover") as gen_cover:
|
||||
summary = ai.generate_batch(
|
||||
[complete_task],
|
||||
{"title": "标题提示", "cover": "封面"},
|
||||
ai_cfg={"config": cfg, "db_path": cfg["db_path"]},
|
||||
)
|
||||
|
||||
self.assertTrue(summary["ok"])
|
||||
self.assertEqual(0, summary["total"])
|
||||
self.assertEqual(0, summary["title_total"])
|
||||
self.assertEqual(0, summary["cover_total"])
|
||||
gen_title.assert_not_called()
|
||||
gen_cover.assert_not_called()
|
||||
updated = db.list_tasks(batch_id=batch_id, path=cfg["db_path"])[0]
|
||||
self.assertEqual("generated", updated.stage)
|
||||
self.assertEqual("success", updated.status)
|
||||
self.assertEqual("重新生成标题", updated.new_title)
|
||||
self.assertIsNone(updated.new_cover_path)
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user