from io import BytesIO from uuid import uuid4 import pytest from PIL import Image from pytest_mock.plugin import MockerFixture from lnbits.core.crud import create_account from lnbits.core.crud.assets import get_user_asset, get_user_assets_count from lnbits.core.models import Account from lnbits.core.services.assets import create_user_asset, thumbnail_from_bytes from lnbits.settings import Settings from tests.helpers import make_upload_file @pytest.mark.anyio async def test_create_user_asset_validates_upload_constraints( app, settings: Settings, mocker: MockerFixture ): file_without_type = make_upload_file(b"hello", filename="a.txt", content_type=None) with pytest.raises(ValueError, match="File must have a content type."): await create_user_asset("user-1", file_without_type, is_public=False) bad_type = make_upload_file( b"hello", filename="bad.bin", content_type="application/x-msdownload", ) with pytest.raises( ValueError, match="File type 'application/x-msdownload' not allowed." ): await create_user_asset("user-1", bad_type, is_public=False) xsl_upload = make_upload_file( b"hello", filename="style.xsl", content_type="text/xml", ) with pytest.raises(ValueError, match="File type 'text/xml' not allowed."): await create_user_asset("user-1", xsl_upload, is_public=False) original_allowed_mime_types = list(settings.lnbits_assets_allowed_mime_types) try: settings.lnbits_assets_allowed_mime_types = [ *original_allowed_mime_types, "text/xml", ] xsl_content = make_upload_file( b'', filename="style.xml", content_type="text/xml", ) with pytest.raises(ValueError, match="File type 'text/xml' not allowed."): await create_user_asset("user-1", xsl_content, is_public=False) finally: settings.lnbits_assets_allowed_mime_types = original_allowed_mime_types fake_image = make_upload_file( b"", filename="fake.png", content_type="image/png", ) with pytest.raises( ValueError, match="Image file content does not match declared file type." ): await create_user_asset("user-1", fake_image, is_public=False) original_max_assets = settings.lnbits_max_assets_per_user original_max_size = settings.lnbits_max_asset_size_mb original_no_limit_users = list(settings.lnbits_assets_no_limit_users) try: settings.lnbits_max_assets_per_user = 1 settings.lnbits_max_asset_size_mb = 1 settings.lnbits_assets_no_limit_users = [] limited_user = await _create_user() allowed_type = make_upload_file( _png_bytes(), filename="ok.png", content_type="image/png" ) await create_user_asset(limited_user, allowed_type, is_public=False) blocked_by_count = make_upload_file( _png_bytes(), filename="again.png", content_type="image/png", ) with pytest.raises(ValueError, match="Max upload count of 1 exceeded."): await create_user_asset(limited_user, blocked_by_count, is_public=False) settings.lnbits_max_asset_size_mb = 0.000001 oversized_user = await _create_user() large_file = make_upload_file( _png_bytes(), filename="ok.png", content_type="image/png", ) with pytest.raises(ValueError, match="File limit of 1e-06MB exceeded."): await create_user_asset(oversized_user, large_file, is_public=False) finally: settings.lnbits_max_assets_per_user = original_max_assets settings.lnbits_max_asset_size_mb = original_max_size settings.lnbits_assets_no_limit_users = original_no_limit_users @pytest.mark.anyio async def test_create_user_asset_success(app, mocker: MockerFixture): user_id = await _create_user() mocker.patch( "lnbits.core.services.assets.thumbnail_from_bytes", return_value=None, ) contents = _png_bytes() file = make_upload_file(contents, filename="hello.png", content_type="image/png") asset = await create_user_asset(user_id, file, is_public=True) stored = await get_user_asset(user_id, asset.id) assert asset.id assert asset.user_id == user_id assert asset.name == "hello.png" assert asset.size_bytes == len(contents) assert asset.data == contents assert asset.is_public is True assert stored is not None assert stored.id == asset.id assert stored.data == contents assert await get_user_assets_count(user_id) == 1 @pytest.mark.anyio async def test_create_user_asset_stores_detected_image_mime_type(app): user_id = await _create_user() buffer = BytesIO() Image.new("RGB", (32, 32), color="blue").save(buffer, format="JPEG") file = make_upload_file( buffer.getvalue(), filename="photo.jpg", content_type="image/jpg" ) asset = await create_user_asset(user_id, file, is_public=True) stored = await get_user_asset(user_id, asset.id) assert asset.mime_type == "image/jpeg" assert stored is not None assert stored.mime_type == "image/jpeg" @pytest.mark.anyio async def test_create_user_asset_rejects_mismatched_image_content(app): user_id = await _create_user() buffer = BytesIO() Image.new("RGB", (32, 32), color="blue").save(buffer, format="JPEG") file = make_upload_file( buffer.getvalue(), filename="photo.png", content_type="image/png" ) with pytest.raises( ValueError, match=( "Image file content does not match declared file type. " "Declared: 'image/png', detected: 'image/jpeg'." ), ): await create_user_asset(user_id, file, is_public=False) def test_thumbnail_from_bytes_success_and_failure(): image = Image.new("RGB", (512, 512), color="red") buffer = BytesIO() image.save(buffer, format="PNG") thumbnail = thumbnail_from_bytes(buffer.getvalue()) assert thumbnail is not None assert isinstance(thumbnail.getvalue(), bytes) assert thumbnail_from_bytes(b"not-an-image") is None async def _create_user() -> str: user_id = uuid4().hex await create_account(Account(id=user_id, username=f"user_{user_id[:8]}")) return user_id def _png_bytes() -> bytes: buffer = BytesIO() Image.new("RGB", (32, 32), color="green").save(buffer, format="PNG") return buffer.getvalue()