import os import urllib.parse from pathlib import Path from unittest.mock import patch import pytest from aiohttp import web from aiohttp.test_utils import make_mocked_request from py.config import Config from py.routes.handlers.preview_handlers import PreviewHandler async def test_preview_handler_serves_preview_from_active_library(tmp_path): library_root = tmp_path / "library" library_root.mkdir() preview_file = library_root / "model.webp" preview_file.write_bytes(b"preview") config = Config() config.apply_library_settings( { "folder_paths": { "loras": [str(library_root)], "checkpoints": [], "unet": [], "embeddings": [], } } ) handler = PreviewHandler(config=config) encoded_path = urllib.parse.quote(str(preview_file), safe="") request = make_mocked_request("GET", f"/api/lm/previews?path={encoded_path}") response = await handler.serve_preview(request) assert isinstance(response, web.FileResponse) assert response.status == 200 assert Path(response._path) == preview_file # TODO: disable temporarily. Enable this once the symlink scan bug fixed # async def test_preview_handler_forbids_paths_outside_active_library(tmp_path): # allowed_root = tmp_path / "allowed" # allowed_root.mkdir() # forbidden_root = tmp_path / "forbidden" # forbidden_root.mkdir() # forbidden_file = forbidden_root / "sneaky.webp" # forbidden_file.write_bytes(b"x") # config = Config() # config.apply_library_settings( # { # "folder_paths": { # "loras": [str(allowed_root)], # "checkpoints": [], # "unet": [], # "embeddings": [], # } # } # ) # handler = PreviewHandler(config=config) # encoded_path = urllib.parse.quote(str(forbidden_file), safe="") # request = make_mocked_request("GET", f"/api/lm/previews?path={encoded_path}") # with pytest.raises(web.HTTPForbidden): # await handler.serve_preview(request) async def test_config_updates_preview_roots_after_switch(tmp_path): first_root = tmp_path / "first" first_root.mkdir() second_root = tmp_path / "second" second_root.mkdir() first_preview = first_root / "model.webp" first_preview.write_bytes(b"a") second_preview = second_root / "model.webp" second_preview.write_bytes(b"b") config = Config() config.apply_library_settings( { "folder_paths": { "loras": [str(first_root)], "checkpoints": [], "unet": [], "embeddings": [], } } ) assert config.is_preview_path_allowed(str(first_preview)) assert not config.is_preview_path_allowed(str(second_preview)) config.apply_library_settings( { "folder_paths": { "loras": [str(second_root)], "checkpoints": [], "unet": [], "embeddings": [], } } ) assert config.is_preview_path_allowed(str(second_preview)) assert not config.is_preview_path_allowed(str(first_preview)) preview_url = config.get_preview_static_url(str(second_preview)) assert preview_url.startswith("/api/lm/previews?path=") decoded = urllib.parse.unquote(preview_url.split("path=", 1)[1]) assert decoded.replace("\\", "/").endswith("model.webp") def test_is_preview_path_allowed_case_insensitive_on_windows(tmp_path): """Test that preview path validation is case-insensitive on Windows. On Windows, drive letters and paths are case-insensitive. This test verifies that paths like 'a:/folder/file' match roots stored as 'A:/folder'. See: https://github.com/willmiao/ComfyUI-Lora-Manager/issues/772 See: https://github.com/willmiao/ComfyUI-Lora-Manager/issues/774 """ # Create actual files for the test library_root = tmp_path / "loras" library_root.mkdir() preview_file = library_root / "model.preview.jpeg" preview_file.write_bytes(b"preview") config = Config() # Simulate Windows behavior by mocking os.path.normcase to lowercase paths # and os.sep to backslash, regardless of the actual platform def windows_normcase(path): return path.lower().replace("/", "\\") with patch("py.config.os.path.normcase", side_effect=windows_normcase), \ patch("py.config.os.sep", "\\"): # Manually set _preview_root_paths with uppercase drive letter style path uppercase_root = Path(str(library_root).upper()) config._preview_root_paths = {uppercase_root} # Test: lowercase version of the path should still be allowed lowercase_path = str(preview_file).lower() assert config.is_preview_path_allowed(lowercase_path), \ f"Path '{lowercase_path}' should be allowed when root is '{uppercase_root}'" # Test: mixed case should also work mixed_case_path = str(preview_file).swapcase() assert config.is_preview_path_allowed(mixed_case_path), \ f"Path '{mixed_case_path}' should be allowed when root is '{uppercase_root}'" # Test: path outside root should still be rejected outside_path = str(tmp_path / "other" / "file.jpeg") assert not config.is_preview_path_allowed(outside_path), \ f"Path '{outside_path}' should NOT be allowed" def test_is_preview_path_allowed_rejects_prefix_without_separator(tmp_path): """Test that 'A:/folder' does not match 'A:/folderextra/file'. This ensures we check for the path separator after the root to avoid false positives with paths that share a common prefix. """ library_root = tmp_path / "loras" library_root.mkdir() # Create a sibling folder that starts with the same prefix sibling_root = tmp_path / "loras_extra" sibling_root.mkdir() sibling_file = sibling_root / "model.jpeg" sibling_file.write_bytes(b"x") config = Config() config.apply_library_settings( { "folder_paths": { "loras": [str(library_root)], "checkpoints": [], "unet": [], "embeddings": [], } } ) # The sibling path should NOT be allowed even though it shares a prefix assert not config.is_preview_path_allowed(str(sibling_file)), \ f"Path in '{sibling_root}' should NOT be allowed when root is '{library_root}'"