mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-10-07 18:12:12 -03:00
Will Miao
470d85cca6
Refreshing had no way to say "scan only this drive": a user with three external
drives had to spin all of them up for every refresh, and switching a drive off
made the next refresh treat its whole library as deleted (rows pruned from the
memory cache and the SQLite cache, preview_url stripped on the next scroll).
Backend (py/services/model_scanner.py, py/config.py):
* ReconcileScope(roots, folder) + _reconcile_cache(scope=...): files inside the
scope reconcile normally, everything outside is neither re-read nor removed.
The folder half is plumbing for the sidebar entry in the next change.
* Path-level pruning guard: cached entries under a path this walk could not
read are kept and reported instead of removed. Sources: a configured root that
is not reachable (drive switched off while LM runs), a directory os.walk
failed to enter (permissions / I/O error / Windows junction to an offline
drive), and a known first-level symlink whose target is gone
(Config.iter_path_mappings()).
* The recorded folder list is unioned instead of replaced whenever the scan did
not verify every root, so a scoped scan cannot empty the sidebar.
* _reconcile_cache returns a summary (added / removed / repaired /
scanned_roots / skipped_roots / unavailable_paths / kept_unreachable),
exposed as ModelScanner.last_reconcile_summary, returned by
BaseModelService.scan_models() and broadcast in the completed WS payload.
* _root_display_labels(): set-aware labels ("G: loras", "usb/loras") grown
leftwards with real parent segments until unique, shared by the walk-progress
line and the roots API.
* GET /scan accepts repeated `roots` (400 for unknown roots, 400 combined with
full_rebuild=true); GET /roots gains root_details (label / reachable / cached
count) while `roots` stays a plain path list for existing callers.
* serve_preview: a 404 no longer clears the cached preview_url when the file's
own directory is unreachable - browsing the grid with a drive off used to
strip preview references from the persistent cache.
Frontend:
* Refresh ▾ gains a "Scan one folder" section listing the page's roots with
their cached counts; offline roots stay clickable and explain themselves; rows
are wired by delegation (new static/js/components/controls/ScanScopeMenu.js).
* A scoped scan reports "Scanned <root>: N new, M removed"; a scan that kept
entries reports "<N> models kept: <paths> not reachable".
* registerAPI() now injects the two cross-page passthroughs (fetchModelRoots and
an argument-forwarding refreshModels) so a page facade cannot drop them: the
first version rendered an empty menu and would have run a full refresh.
* createToastElement whitelists toast types, so a wrong `type` argument degrades
to the info style instead of rendering an unstyled box.
Verified in a sandbox instance with three roots: a scoped scan walks only the
requested root (progress roots=0/1, 240 files); a full refresh with one root
offline reports kept_unreachable=60 and leaves all 420 models cached; /roots
reports the offline root with its cached count. 3699 passed, 7 skipped;
frontend 1488 passed (148 files); vue widgets 96 passed.
416 lines
14 KiB
Python
416 lines
14 KiB
Python
import asyncio
|
|
import os
|
|
import urllib.parse
|
|
from pathlib import Path
|
|
from unittest.mock import AsyncMock, 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
|
|
|
|
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")
|
|
|
|
|
|
async def test_preview_handler_allows_custom_recipes_path(tmp_path):
|
|
lora_root = tmp_path / "library"
|
|
lora_root.mkdir()
|
|
recipes_root = tmp_path / "recipes_storage"
|
|
recipes_root.mkdir()
|
|
preview_file = recipes_root / "recipe.webp"
|
|
preview_file.write_bytes(b"preview")
|
|
|
|
config = Config()
|
|
config.apply_library_settings(
|
|
{
|
|
"folder_paths": {
|
|
"loras": [str(lora_root)],
|
|
"checkpoints": [],
|
|
"unet": [],
|
|
"embeddings": [],
|
|
},
|
|
"recipes_path": str(recipes_root),
|
|
}
|
|
)
|
|
|
|
assert config.is_preview_path_allowed(str(preview_file))
|
|
|
|
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
|
|
|
|
|
|
async def test_preview_handler_allows_symlinked_recipes_path(tmp_path):
|
|
lora_root = tmp_path / "library"
|
|
lora_root.mkdir()
|
|
real_recipes_root = tmp_path / "real_recipes"
|
|
real_recipes_root.mkdir()
|
|
symlink_recipes_root = tmp_path / "linked_recipes"
|
|
symlink_recipes_root.symlink_to(real_recipes_root, target_is_directory=True)
|
|
|
|
preview_file = real_recipes_root / "recipe.webp"
|
|
preview_file.write_bytes(b"preview")
|
|
|
|
config = Config()
|
|
config.apply_library_settings(
|
|
{
|
|
"folder_paths": {
|
|
"loras": [str(lora_root)],
|
|
"checkpoints": [],
|
|
"unet": [],
|
|
"embeddings": [],
|
|
},
|
|
"recipes_path": str(symlink_recipes_root),
|
|
}
|
|
)
|
|
|
|
symlink_preview_path = symlink_recipes_root / "recipe.webp"
|
|
assert config.is_preview_path_allowed(str(symlink_preview_path))
|
|
|
|
handler = PreviewHandler(config=config)
|
|
encoded_path = urllib.parse.quote(str(symlink_preview_path), 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.resolve()
|
|
|
|
|
|
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}'"
|
|
|
|
|
|
async def test_preview_handler_serves_from_deep_symlink(tmp_path):
|
|
"""Test that previews under deep symlinks are served correctly."""
|
|
library_root = tmp_path / "library"
|
|
library_root.mkdir()
|
|
|
|
# Create nested structure with deep symlink at second level
|
|
subdir = library_root / "anime"
|
|
subdir.mkdir()
|
|
external_dir = tmp_path / "external"
|
|
external_dir.mkdir()
|
|
deep_symlink = subdir / "styles"
|
|
deep_symlink.symlink_to(external_dir, target_is_directory=True)
|
|
|
|
# Create preview file under deep symlink
|
|
preview_file = deep_symlink / "model.preview.webp"
|
|
preview_file.write_bytes(b"preview_content")
|
|
|
|
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.resolve()
|
|
|
|
|
|
async def test_deep_symlink_discovered_on_first_access(tmp_path):
|
|
"""Test that deep symlinks are discovered on first preview access."""
|
|
library_root = tmp_path / "library"
|
|
library_root.mkdir()
|
|
|
|
# Create nested structure with deep symlink at second level
|
|
subdir = library_root / "category"
|
|
subdir.mkdir()
|
|
external_dir = tmp_path / "storage"
|
|
external_dir.mkdir()
|
|
deep_symlink = subdir / "models"
|
|
deep_symlink.symlink_to(external_dir, target_is_directory=True)
|
|
|
|
# Create preview file under deep symlink
|
|
preview_file = deep_symlink / "test.png"
|
|
preview_file.write_bytes(b"test_image")
|
|
|
|
config = Config()
|
|
config.apply_library_settings(
|
|
{
|
|
"folder_paths": {
|
|
"loras": [str(library_root)],
|
|
"checkpoints": [],
|
|
"unet": [],
|
|
"embeddings": [],
|
|
}
|
|
}
|
|
)
|
|
|
|
# Deep symlink should not be in mappings initially
|
|
normalized_external = os.path.normpath(str(external_dir)).replace(os.sep, '/')
|
|
assert normalized_external not in config._path_mappings
|
|
|
|
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}")
|
|
|
|
# First access should trigger symlink discovery and serve the preview
|
|
response = await handler.serve_preview(request)
|
|
|
|
assert isinstance(response, web.FileResponse)
|
|
assert response.status == 200
|
|
assert Path(response._path) == preview_file.resolve()
|
|
|
|
# Deep symlink should now be in mappings
|
|
assert normalized_external in config._path_mappings
|
|
|
|
|
|
async def test_preview_404_keeps_cache_when_parent_dir_missing(tmp_path):
|
|
"""A preview whose directory is unreachable (drive off) must not clear the
|
|
cached preview_url: nothing was deleted, the file just cannot be read."""
|
|
library_root = tmp_path / "library"
|
|
library_root.mkdir()
|
|
unreachable_parent = library_root / "gone"
|
|
preview_path = unreachable_parent / "model.webp"
|
|
|
|
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_path), safe="")
|
|
request = make_mocked_request("GET", f"/api/lm/previews?path={encoded_path}")
|
|
|
|
with patch.object(handler, "_cleanup_stale_preview_url", new=AsyncMock()) as cleanup:
|
|
with pytest.raises(web.HTTPNotFound):
|
|
await handler.serve_preview(request)
|
|
await asyncio.sleep(0)
|
|
|
|
cleanup.assert_not_called()
|
|
|
|
|
|
async def test_preview_404_clears_cache_when_file_really_deleted(tmp_path):
|
|
"""A preview deleted from a reachable directory is genuinely stale."""
|
|
library_root = tmp_path / "library"
|
|
library_root.mkdir()
|
|
preview_path = library_root / "model.webp"
|
|
preview_path.write_bytes(b"preview")
|
|
preview_path.unlink()
|
|
|
|
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_path), safe="")
|
|
request = make_mocked_request("GET", f"/api/lm/previews?path={encoded_path}")
|
|
|
|
with patch.object(handler, "_cleanup_stale_preview_url", new=AsyncMock()) as cleanup:
|
|
with pytest.raises(web.HTTPNotFound):
|
|
await handler.serve_preview(request)
|
|
await asyncio.sleep(0)
|
|
|
|
cleanup.assert_awaited_once()
|