ComfyUI-Lora-Manager/tests/routes/test_preview_routes.py

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}'"