279 lines
10 KiB
Python
279 lines
10 KiB
Python
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_generation_round_copies_successes_and_uses_new_directory_on_repeat(self):
|
|
with self.make_temp_dir() as temp_dir:
|
|
cfg, project, _main, _detail = self._project_with_assets(temp_dir)
|
|
source = image_studio.add_asset(
|
|
project.id,
|
|
image_studio.ASSET_KIND_ORIGINAL,
|
|
path=cfg["db_path"],
|
|
)
|
|
first_source = self._write_image(os.path.join(temp_dir, "round-first.png"))
|
|
first_asset = image_studio.add_asset(
|
|
project.id,
|
|
"generated_main",
|
|
local_path=first_source,
|
|
parent_asset_id=source.id,
|
|
path=cfg["db_path"],
|
|
)
|
|
first_job = image_studio.create_job(
|
|
project.id,
|
|
source_asset_id=source.id,
|
|
job_type="白底图",
|
|
task_key="round-export-first",
|
|
generation_round_key="round-export",
|
|
generation_slot_index=0,
|
|
path=cfg["db_path"],
|
|
)
|
|
image_studio.update_job_status(
|
|
first_job.id,
|
|
"succeeded",
|
|
output_asset_id=first_asset.id,
|
|
path=cfg["db_path"],
|
|
)
|
|
missing_asset = image_studio.add_asset(
|
|
project.id,
|
|
"generated_main",
|
|
local_path=os.path.join(temp_dir, "missing-round.png"),
|
|
parent_asset_id=source.id,
|
|
path=cfg["db_path"],
|
|
)
|
|
missing_job = image_studio.create_job(
|
|
project.id,
|
|
source_asset_id=source.id,
|
|
job_type="场景图",
|
|
task_key="round-export-missing",
|
|
generation_round_key="round-export",
|
|
generation_slot_index=1,
|
|
path=cfg["db_path"],
|
|
)
|
|
image_studio.update_job_status(
|
|
missing_job.id,
|
|
"succeeded",
|
|
output_asset_id=missing_asset.id,
|
|
path=cfg["db_path"],
|
|
)
|
|
second_source = self._write_image(
|
|
os.path.join(temp_dir, "round-second.png"),
|
|
color=(220, 80, 40, 255),
|
|
)
|
|
second_asset = image_studio.add_asset(
|
|
project.id,
|
|
"generated_main",
|
|
local_path=second_source,
|
|
parent_asset_id=source.id,
|
|
path=cfg["db_path"],
|
|
)
|
|
second_job = image_studio.create_job(
|
|
project.id,
|
|
source_asset_id=source.id,
|
|
job_type="卖点图",
|
|
task_key="round-export-second",
|
|
generation_round_key="round-export",
|
|
generation_slot_index=2,
|
|
path=cfg["db_path"],
|
|
)
|
|
image_studio.update_job_status(
|
|
second_job.id,
|
|
"succeeded",
|
|
output_asset_id=second_asset.id,
|
|
path=cfg["db_path"],
|
|
)
|
|
parent = os.path.join(temp_dir, "exports")
|
|
os.makedirs(parent)
|
|
|
|
result = image_studio_export.export_generation_round(
|
|
project.id,
|
|
"round-export",
|
|
parent,
|
|
path=cfg["db_path"],
|
|
)
|
|
|
|
self.assertEqual(2, len(result.files))
|
|
self.assertEqual(1, result.skipped_count)
|
|
self.assertFalse(result.cancelled)
|
|
self.assertTrue(os.path.isdir(result.target_dir))
|
|
self.assertEqual(
|
|
["01_白底图.png", "03_卖点图.png"],
|
|
sorted(os.path.basename(item.output_path) for item in result.files),
|
|
)
|
|
self.assertTrue(os.path.isfile(first_source))
|
|
self.assertTrue(os.path.isfile(second_source))
|
|
with open(first_source, "rb") as source_file, open(
|
|
result.files[0].output_path,
|
|
"rb",
|
|
) as exported_file:
|
|
self.assertEqual(source_file.read(), exported_file.read())
|
|
|
|
second_result = image_studio_export.export_generation_round(
|
|
project.id,
|
|
"round-export",
|
|
parent,
|
|
path=cfg["db_path"],
|
|
)
|
|
self.assertNotEqual(result.target_dir, second_result.target_dir)
|
|
self.assertTrue(second_result.target_dir.endswith("_2"))
|
|
|
|
self.assert_removed(temp_dir)
|
|
|
|
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()
|