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