cmshoppe/tests/test_image_studio_export.py

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