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()