468 lines
18 KiB
Python
468 lines
18 KiB
Python
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()
|