diff --git a/tests-unit/app_test/model_manager_test.py b/tests-unit/app_test/model_manager_test.py new file mode 100644 index 000000000..ae59206f6 --- /dev/null +++ b/tests-unit/app_test/model_manager_test.py @@ -0,0 +1,62 @@ +import pytest +import base64 +import json +import struct +from io import BytesIO +from PIL import Image +from aiohttp import web +from unittest.mock import patch +from app.model_manager import ModelFileManager + +pytestmark = ( + pytest.mark.asyncio +) # This applies the asyncio mark to all test functions in the module + +@pytest.fixture +def model_manager(): + return ModelFileManager() + +@pytest.fixture +def app(model_manager): + app = web.Application() + routes = web.RouteTableDef() + model_manager.add_routes(routes) + app.add_routes(routes) + return app + +async def test_get_model_preview_safetensors(aiohttp_client, app, tmp_path): + img = Image.new('RGB', (100, 100), 'white') + img_byte_arr = BytesIO() + img.save(img_byte_arr, format='PNG') + img_byte_arr.seek(0) + img_b64 = base64.b64encode(img_byte_arr.getvalue()).decode('utf-8') + + safetensors_file = tmp_path / "test_model.safetensors" + header_bytes = json.dumps({ + "__metadata__": { + "ssmd_cover_images": json.dumps([img_b64]) + } + }).encode('utf-8') + length_bytes = struct.pack('