cmshoppe/tests/test_image_studio_images.py

import io
import os
import socket
import sys
import threading
import time
import unittest
from types import SimpleNamespace
from unittest import mock

sys.path.insert(0, os.path.dirname(__file__))

from _helpers import TempDirMixin

from app import db, image_studio, image_studio_images


class _Response:
    def __init__(
        self,
        content=b"",
        status_code=200,
        headers=None,
        url="https://cdn.example.com/image.png",
        chunk_size=None,
    ):
        self.content = content
        self.status_code = status_code
        self.headers = dict(headers or {})
        self.url = url
        self.chunk_size = chunk_size
        self.closed = False

    def iter_content(self, chunk_size=65536):
        if self.chunk_size:
            for index in range(0, len(self.content), self.chunk_size):
                yield self.content[index:index + self.chunk_size]
            return
        yield self.content

    def close(self):
        self.closed = True


class ImageStudioImageTests(TempDirMixin, unittest.TestCase):
    def _png_bytes(self, size=(20, 16), color=(80, 120, 200, 255)):
        from PIL import Image

        output = io.BytesIO()
        Image.new("RGBA", size, color).save(output, format="PNG")
        return output.getvalue()

    def _public_dns(self):
        return [(socket.AF_INET, socket.SOCK_STREAM, 6, "", ("93.184.216.34", 443))]

    def _project_and_asset(self, temp_dir, remote_url="https://cdn.example.com/source.png"):
        db_path = os.path.join(temp_dir, "cmshopee.db")
        db.init_db(db_path)
        project = image_studio.create_or_get_project(
            account_alias="alias",
            account_slug="alias_slug",
            item_id="51100639510",
            path=db_path,
        )
        asset = image_studio.sync_original_asset_urls(
            project.id,
            [{"index": 1, "src": remote_url}],
            path=db_path,
        )[0]
        return db_path, project, asset

    def test_download_remote_image_rejects_private_non_image_overlimit_and_private_redirect(self):
        with self.assertRaisesRegex(image_studio_images.ImageStudioImageError, "http/https|本机或内网"):
            image_studio_images.download_remote_image("http://127.0.0.1/a.png")

        non_image_session = SimpleNamespace(
            get=lambda *args, **kwargs: _Response(
                content=b"hello",
                headers={"Content-Type": "text/html"},
            )
        )
        with mock.patch("app.image_studio_images.socket.getaddrinfo", return_value=self._public_dns()):
            with self.assertRaisesRegex(image_studio_images.ImageStudioImageError, "不是图片"):
                image_studio_images.download_remote_image(
                    "https://cdn.example.com/a.png",
                    session=non_image_session,
                )

        oversized_session = SimpleNamespace(
            get=lambda *args, **kwargs: _Response(
                content=b"x",
                headers={"Content-Type": "image/png", "Content-Length": "999"},
            )
        )
        with mock.patch("app.image_studio_images.socket.getaddrinfo", return_value=self._public_dns()):
            with self.assertRaisesRegex(image_studio_images.ImageStudioImageError, "大小上限"):
                image_studio_images.download_remote_image(
                    "https://cdn.example.com/a.png",
                    session=oversized_session,
                    max_bytes=10,
                )

        redirect_session = SimpleNamespace(
            get=lambda *args, **kwargs: _Response(
                content=self._png_bytes(),
                headers={"Content-Type": "image/png"},
                url="http://127.0.0.1/private.png",
            )
        )
        with mock.patch("app.image_studio_images.socket.getaddrinfo", return_value=self._public_dns()):
            with self.assertRaisesRegex(image_studio_images.ImageStudioImageError, "本机或内网"):
                image_studio_images.download_remote_image(
                    "https://cdn.example.com/a.png",
                    session=redirect_session,
                )

    def test_load_thumbnail_decodes_image_without_writing_file(self):
        png = self._png_bytes(size=(640, 480))
        session = SimpleNamespace(
            get=lambda *args, **kwargs: _Response(
                content=png,
                headers={"Content-Type": "image/png"},
            )
        )

        with mock.patch("app.image_studio_images.socket.getaddrinfo", return_value=self._public_dns()):
            result = image_studio_images.load_thumbnail(
                "https://cdn.example.com/a.png",
                key="thumb-1",
                max_size=120,
                session=session,
            )

        self.assertEqual("thumb-1", result.key)
        self.assertLessEqual(result.width, 120)
        self.assertLessEqual(result.height, 120)
        self.assertTrue(result.image_bytes.startswith(b"\x89PNG"))

    def test_thumbnail_loader_uses_memory_cache_and_can_cancel(self):
        calls = []
        successes = []
        errors = []

        def fake_loader(url, key=None):
            calls.append((key, url))
            return image_studio_images.ThumbnailResult(
                key=str(key),
                url=url,
                image_bytes=b"png",
                width=10,
                height=10,
            )

        loader = image_studio_images.ThumbnailLoader(max_workers=1, loader=fake_loader)
        try:
            future = loader.submit("a", "https://cdn.example.com/a.png", successes.append, errors.append)
            future.result(timeout=2)
            for _ in range(20):
                if successes:
                    break
                time.sleep(0.01)
            loader.submit("b", "https://cdn.example.com/a.png", successes.append, errors.append)

            self.assertEqual([("a", "https://cdn.example.com/a.png")], calls)
            self.assertEqual(2, len(successes))
            self.assertFalse(successes[0].from_cache)
            self.assertTrue(successes[1].from_cache)
            self.assertEqual([], errors)

            release = threading.Event()

            def slow_loader(url, key=None):
                release.wait(1)
                return image_studio_images.ThumbnailResult(
                    key=str(key),
                    url=url,
                    image_bytes=b"png",
                    width=10,
                    height=10,
                )

            loader._loader = slow_loader
            future = loader.submit("c", "https://cdn.example.com/c.png", successes.append, errors.append)
            loader.cancel("c")
            release.set()
            future.result(timeout=2)
            time.sleep(0.05)
            self.assertEqual(2, len(successes))
        finally:
            loader.close()

    def test_download_original_asset_writes_atomically_and_is_idempotent(self):
        with self.make_temp_dir() as temp_dir:
            db_path, _, asset = self._project_and_asset(temp_dir)
            png = self._png_bytes()
            session = mock.Mock()
            session.get.return_value = _Response(
                content=png,
                headers={"Content-Type": "image/png"},
            )

            with mock.patch("app.image_studio_images.socket.getaddrinfo", return_value=self._public_dns()):
                downloaded = image_studio_images.download_original_asset(
                    asset.id,
                    path=db_path,
                    config={"data_dir": temp_dir, "image_dir": os.path.join(temp_dir, "images")},
                    session=session,
                )
                second = image_studio_images.download_original_asset(
                    asset.id,
                    path=db_path,
                    config={"data_dir": temp_dir, "image_dir": os.path.join(temp_dir, "images")},
                    session=mock.Mock(),
                )

            self.assertEqual(downloaded.id, second.id)
            self.assertEqual(downloaded.local_path, second.local_path)
            self.assertTrue(os.path.isfile(downloaded.local_path))
            self.assertTrue(downloaded.local_path.endswith(".png"))
            self.assertEqual(image_studio.ASSET_STATUS_AVAILABLE, downloaded.status)
            session.get.assert_called_once()
            originals = os.path.dirname(downloaded.local_path)
            self.assertFalse([name for name in os.listdir(originals) if ".tmp-" in name])

        self.assert_removed(temp_dir)

    def test_download_original_asset_bad_image_leaves_asset_without_local_path(self):
        with self.make_temp_dir() as temp_dir:
            db_path, project, asset = self._project_and_asset(temp_dir)
            session = SimpleNamespace(
                get=lambda *args, **kwargs: _Response(
                    content=b"not an image",
                    headers={"Content-Type": "image/png"},
                )
            )

            with mock.patch("app.image_studio_images.socket.getaddrinfo", return_value=self._public_dns()):
                with self.assertRaisesRegex(image_studio_images.ImageStudioImageError, "解码失败"):
                    image_studio_images.download_original_asset(
                        asset.id,
                        path=db_path,
                        config={"data_dir": temp_dir, "image_dir": os.path.join(temp_dir, "images")},
                        session=session,
                    )

            refreshed = image_studio.get_asset(asset.id, path=db_path)
            self.assertIsNone(refreshed.local_path)
            dirs = image_studio.project_image_dirs(os.path.join(temp_dir, "images"), project)
            self.assertFalse(os.path.exists(dirs["originals"]))

        self.assert_removed(temp_dir)

    def test_download_remote_image_cancels_between_chunks_and_closes_response(self):
        png = self._png_bytes(size=(120, 120))
        response = _Response(
            content=png,
            headers={"Content-Type": "image/png"},
            chunk_size=32,
        )
        session = SimpleNamespace(get=lambda *args, **kwargs: response)
        checks = {"count": 0}

        def should_stop():
            checks["count"] += 1
            return checks["count"] >= 4

        with mock.patch(
            "app.image_studio_images.socket.getaddrinfo",
            return_value=self._public_dns(),
        ):
            with self.assertRaises(image_studio_images.ImageStudioImageCancelled):
                image_studio_images.download_remote_image(
                    "https://cdn.example.com/a.png",
                    session=session,
                    should_stop=should_stop,
                )

        self.assertTrue(response.closed)

    def test_import_original_files_copies_valid_images_deduplicates_and_limits(self):
        with self.make_temp_dir() as temp_dir:
            db_path = os.path.join(temp_dir, "cmshopee.db")
            db.init_db(db_path)
            project = image_studio.create_or_get_project(
                account_alias="alias",
                account_slug="alias_slug",
                item_id="51100639510",
                path=db_path,
            )
            source = os.path.join(temp_dir, "商品图.png")
            with open(source, "wb") as fh:
                fh.write(self._png_bytes())
            config = {
                "data_dir": temp_dir,
                "db_path": db_path,
                "image_dir": os.path.join(temp_dir, "images"),
            }

            first = image_studio_images.import_original_files(
                project.id,
                [source],
                path=db_path,
                config=config,
            )
            duplicate = image_studio_images.import_original_files(
                project.id,
                [source],
                path=db_path,
                config=config,
            )

            self.assertEqual([], first["errors"])
            self.assertEqual(first["assets"][0].id, duplicate["assets"][0].id)
            self.assertNotEqual(os.path.abspath(source), first["assets"][0].local_path)
            self.assertTrue(os.path.isfile(first["assets"][0].local_path))
            with self.assertRaisesRegex(image_studio_images.ImageStudioImageError, "最多"):
                image_studio_images.import_original_bytes(
                    project.id,
                    self._png_bytes(color=(10, 20, 30, 255)),
                    filename_hint="second.png",
                    path=db_path,
                    config=config,
                    max_assets=1,
                )

        self.assert_removed(temp_dir)

    def test_draft_import_uses_stable_storage_key_after_formal_binding(self):
        with self.make_temp_dir() as temp_dir:
            db_path = os.path.join(temp_dir, "cmshopee.db")
            db.init_db(db_path)
            draft = image_studio.create_draft_project(
                account_alias="alias",
                account_slug="alias_slug",
                path=db_path,
            )
            config = {
                "data_dir": temp_dir,
                "db_path": db_path,
                "image_dir": os.path.join(temp_dir, "images"),
            }

            first = image_studio_images.import_original_bytes(
                draft.id,
                self._png_bytes(color=(10, 20, 30, 255)),
                filename_hint="first.png",
                path=db_path,
                config=config,
            )
            first_path = first.local_path
            bound = image_studio.bind_draft_project(draft.id, "51100639510", path=db_path)
            second = image_studio_images.import_original_bytes(
                bound.id,
                self._png_bytes(color=(40, 50, 60, 255)),
                filename_hint="second.png",
                path=db_path,
                config=config,
            )

            self.assertEqual(image_studio.PROJECT_BINDING_BOUND, bound.binding_state)
            self.assertEqual(draft.storage_key, bound.storage_key)
            self.assertEqual(os.path.dirname(first_path), os.path.dirname(second.local_path))
            self.assertTrue(os.path.isfile(first_path))
            self.assertTrue(os.path.isfile(second.local_path))
            self.assertIn(draft.storage_key, first_path)
            self.assertNotIn("51100639510", first_path)

        self.assert_removed(temp_dir)

    def test_import_original_ignores_missing_history_when_enforcing_limit(self):
        with self.make_temp_dir() as temp_dir:
            db_path = os.path.join(temp_dir, "cmshopee.db")
            db.init_db(db_path)
            project = image_studio.create_or_get_project(
                account_alias="alias",
                account_slug="alias_slug",
                item_id="51100639510",
                path=db_path,
            )
            config = {
                "data_dir": temp_dir,
                "db_path": db_path,
                "image_dir": os.path.join(temp_dir, "images"),
            }
            first = image_studio_images.import_original_bytes(
                project.id,
                self._png_bytes(),
                filename_hint="first.png",
                path=db_path,
                config=config,
                max_assets=1,
            )
            image_studio.mark_asset_status(
                first.id,
                image_studio.ASSET_STATUS_MISSING,
                path=db_path,
            )

            second = image_studio_images.import_original_bytes(
                project.id,
                self._png_bytes(color=(10, 20, 30, 255)),
                filename_hint="second.png",
                path=db_path,
                config=config,
                max_assets=1,
            )

            self.assertNotEqual(first.id, second.id)
            self.assertEqual(image_studio.ASSET_STATUS_AVAILABLE, second.status)

        self.assert_removed(temp_dir)

    def test_generated_asset_trash_and_restore_keep_database_history(self):
        with self.make_temp_dir() as temp_dir:
            db_path = os.path.join(temp_dir, "cmshopee.db")
            db.init_db(db_path)
            project = image_studio.create_or_get_project(
                account_alias="alias",
                account_slug="alias_slug",
                item_id="51100639510",
                path=db_path,
            )
            config = {
                "data_dir": temp_dir,
                "db_path": db_path,
                "image_dir": os.path.join(temp_dir, "images"),
            }
            generated_dir = image_studio.default_project_image_dirs(project, config=config)[
                "generated"
            ]
            os.makedirs(generated_dir, exist_ok=True)
            generated_path = os.path.join(generated_dir, "result.png")
            with open(generated_path, "wb") as fh:
                fh.write(self._png_bytes())
            asset = image_studio.add_asset(
                project.id,
                "generated_main",
                local_path=generated_path,
                path=db_path,
            )

            record = image_studio_images.trash_generated_asset(
                asset.id,
                path=db_path,
                config=config,
            )
            trashed = image_studio.get_asset(asset.id, path=db_path)

            self.assertEqual(image_studio.ASSET_STATUS_MISSING, trashed.status)
            self.assertTrue(os.path.isfile(record["trash_path"]))
            self.assertFalse(os.path.exists(generated_path))

            restored = image_studio_images.restore_trashed_asset(
                record,
                path=db_path,
                config=config,
            )

            self.assertEqual(asset.id, restored.id)
            self.assertEqual(image_studio.ASSET_STATUS_AVAILABLE, restored.status)
            self.assertTrue(os.path.isfile(restored.local_path))

        self.assert_removed(temp_dir)


if __name__ == "__main__":
    unittest.main()