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