feat(ai-studio): export final selections
This commit is contained in:
chengma
@@ -0,0 +1,162 @@
|
||||
import io
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
sys.path.insert(0, os.path.dirname(__file__))
|
||||
|
||||
from _helpers import TempDirMixin
|
||||
|
||||
from app import db, image_studio, image_studio_export
|
||||
|
||||
|
||||
class ImageStudioExportTests(TempDirMixin, unittest.TestCase):
|
||||
def _png_bytes(self, color=(20, 120, 200, 180)):
|
||||
from PIL import Image
|
||||
|
||||
output = io.BytesIO()
|
||||
Image.new("RGBA", (24, 24), color).save(output, format="PNG")
|
||||
return output.getvalue()
|
||||
|
||||
def _write_image(self, path, color=(20, 120, 200, 180)):
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
with open(path, "wb") as fh:
|
||||
fh.write(self._png_bytes(color=color))
|
||||
return path
|
||||
|
||||
def _project_with_assets(self, temp_dir):
|
||||
cfg = {
|
||||
"db_path": os.path.join(temp_dir, "cmshopee.db"),
|
||||
"image_dir": os.path.join(temp_dir, "images"),
|
||||
"ai": {"jpg_quality": 88},
|
||||
}
|
||||
db.init_db(cfg["db_path"])
|
||||
project = image_studio.create_or_get_project(
|
||||
account_alias="alias",
|
||||
account_slug="alias_slug",
|
||||
item_id="51100639510",
|
||||
path=cfg["db_path"],
|
||||
)
|
||||
main = image_studio.add_asset(
|
||||
project.id,
|
||||
"generated_main",
|
||||
local_path=self._write_image(os.path.join(temp_dir, "main.png")),
|
||||
path=cfg["db_path"],
|
||||
)
|
||||
detail = image_studio.add_asset(
|
||||
project.id,
|
||||
"generated_detail",
|
||||
local_path=self._write_image(os.path.join(temp_dir, "detail.webp"), color=(200, 80, 40, 255)),
|
||||
path=cfg["db_path"],
|
||||
)
|
||||
image_studio.replace_selections(project.id, "main", [main.id], path=cfg["db_path"])
|
||||
image_studio.replace_selections(project.id, "detail", [detail.id], path=cfg["db_path"])
|
||||
return cfg, project, main, detail
|
||||
|
||||
def test_export_partial_selection_outputs_ordered_jpegs(self):
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
cfg, project, _main, _detail = self._project_with_assets(temp_dir)
|
||||
parent = os.path.join(temp_dir, "exports")
|
||||
os.makedirs(parent)
|
||||
|
||||
result = image_studio_export.export_project_selection(
|
||||
project.id,
|
||||
parent,
|
||||
path=cfg["db_path"],
|
||||
config=cfg,
|
||||
)
|
||||
|
||||
self.assertEqual(1, result.main_count)
|
||||
self.assertEqual(1, result.detail_count)
|
||||
self.assertEqual(
|
||||
["51100639510_主图_1.jpg", "51100639510_详情图_1.jpg"],
|
||||
sorted(os.path.basename(item.output_path) for item in result.files),
|
||||
)
|
||||
from PIL import Image
|
||||
|
||||
for item in result.files:
|
||||
with Image.open(item.output_path) as image:
|
||||
self.assertEqual("JPEG", image.format)
|
||||
self.assertEqual("RGB", image.mode)
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
def test_existing_target_requires_explicit_choice_and_timestamp_mode(self):
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
cfg, project, _main, _detail = self._project_with_assets(temp_dir)
|
||||
parent = os.path.join(temp_dir, "exports")
|
||||
target = os.path.join(parent, project.item_id)
|
||||
os.makedirs(target)
|
||||
|
||||
with self.assertRaises(image_studio_export.ExportTargetExistsError):
|
||||
image_studio_export.export_project_selection(
|
||||
project.id,
|
||||
parent,
|
||||
path=cfg["db_path"],
|
||||
config=cfg,
|
||||
)
|
||||
|
||||
result = image_studio_export.export_project_selection(
|
||||
project.id,
|
||||
parent,
|
||||
existing_mode=image_studio_export.EXISTING_TIMESTAMP,
|
||||
timestamp="20260711_120000",
|
||||
path=cfg["db_path"],
|
||||
config=cfg,
|
||||
)
|
||||
self.assertTrue(result.target_dir.endswith("51100639510_20260711_120000"))
|
||||
self.assertTrue(os.path.isdir(result.target_dir))
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
def test_overwrite_managed_keeps_user_files(self):
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
cfg, project, _main, _detail = self._project_with_assets(temp_dir)
|
||||
parent = os.path.join(temp_dir, "exports")
|
||||
target = os.path.join(parent, project.item_id)
|
||||
os.makedirs(target)
|
||||
managed = os.path.join(target, "51100639510_主图_9.jpg")
|
||||
user_file = os.path.join(target, "用户说明.txt")
|
||||
with open(managed, "w", encoding="utf-8") as fh:
|
||||
fh.write("old")
|
||||
with open(user_file, "w", encoding="utf-8") as fh:
|
||||
fh.write("keep")
|
||||
|
||||
result = image_studio_export.export_project_selection(
|
||||
project.id,
|
||||
parent,
|
||||
existing_mode=image_studio_export.EXISTING_OVERWRITE_MANAGED,
|
||||
path=cfg["db_path"],
|
||||
config=cfg,
|
||||
)
|
||||
|
||||
self.assertFalse(os.path.exists(managed))
|
||||
self.assertTrue(os.path.exists(user_file))
|
||||
with open(user_file, encoding="utf-8") as fh:
|
||||
self.assertEqual("keep", fh.read())
|
||||
self.assertEqual(2, len(result.files))
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
def test_preflight_failure_does_not_create_target(self):
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
cfg, project, main, _detail = self._project_with_assets(temp_dir)
|
||||
os.remove(main.local_path)
|
||||
parent = os.path.join(temp_dir, "exports")
|
||||
os.makedirs(parent)
|
||||
|
||||
with self.assertRaisesRegex(image_studio_export.ImageStudioExportError, "缺失"):
|
||||
image_studio_export.export_project_selection(
|
||||
project.id,
|
||||
parent,
|
||||
path=cfg["db_path"],
|
||||
config=cfg,
|
||||
)
|
||||
|
||||
self.assertFalse(os.path.exists(os.path.join(parent, project.item_id)))
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user