Files
ComfyUI-Lora-Manager/tests/routes/test_model_query_handler.py
T
Will Miao 6d0d7cb15b fix(scanner): compare model roots by identity, not by spelling (#1146)
ComfyUI hands its model folders over in the host's own spelling --
`utils/extra_config.py` runs every configured path through `os.path.normpath()`,
which is all backslashes on Windows, and `main.py` registers `<output>/loras`
unconditionally for its save nodes. `Config` kept that spelling verbatim while
`_prepare_*_paths()` normalized the live root lists to forward slashes, and the
scanner compared the two as raw strings. On Windows every configured root
therefore looked like a root that was not in the live list: the Refresh menu
listed each directory twice (the twin showing "0 models"), the completion toast
claimed reachable roots were "currently inaccessible" and kept their whole cache
count as `kept_unreachable`, and the settings sync persisted host-registered
directories that were never created.

The duplication was not only cosmetic. `_UnreachablePaths.add()` normalizes the
prefix it is given, so the configured twin of a live root registered exactly the
live root's prefix as unreadable: in the prune phase every cached-but-missing
entry under it matched and was "kept", so models deleted from disk were never
removed from the cache again.

* `py/utils/root_paths.py` holds the single `normalize_root_identity()`
  (separator style, redundant/trailing separators, and case where the filesystem
  folds it: whole paths on Windows, drive letters everywhere). `config.py` and
  `settings_manager.py` each carried a private copy and `model_scanner.py`
  compared raw strings.
* `ModelScanner` compares root identities in `get_configured_model_roots()`,
  `describe_model_roots()` (dedupe + `available`), the reconcile's
  `configured_only` list and the `ReconcileScope.roots` membership test, so one
  directory is one root whatever the host called it. `_normalized_root_prefix()`
  is built on the same identity, and an empty prefix no longer matches every
  path.
* A configured root that never made it into the live list is reported only while
  it still holds cached entries -- in the menu and in `skipped_roots` /
  `unavailable_paths` alike. A directory that neither exists nor was ever scanned
  has nothing to keep, so `<output>/loras` stops being an unavailable root.
* The startup sync drops saved roots that neither exist nor hold cached entries
  (`PersistentModelCache.roots_with_cached_entries()`, one scan per model type,
  separator-agnostic and wildcard-escaped). A root named by the library's
  `default_*_root` is kept even when empty: dropping it would let the next boot
  "repair" the default onto another root. When the cache cannot be consulted,
  every root is kept.
* `GET /scan?roots=` validates requested roots by identity too, so the browser
  extension or a script passing the spelling settings.json holds selects the root
  instead of answering 400 "Unknown model root(s)".

Verified: 8 of the new tests fail against the pre-fix source (checked with the
fix stashed); 3798 passed, 7 skipped, 11 deselected. The probe costs 71ms for a
100k-model library with seven missing roots (357ms when probing root by root).
2026-10-12 08:43:06 +08:00

453 lines
15 KiB
Python

import json
import logging
from types import SimpleNamespace
import pytest
from py.routes.handlers.model_handlers import ModelQueryHandler
class DummyService:
def __init__(self):
self.received_limit = None
async def get_base_models(self, limit):
self.received_limit = limit
return [{"name": "SDXL", "count": 2}]
@pytest.mark.asyncio
async def test_model_query_handler_accepts_limit_zero_for_base_models():
service = DummyService()
handler = ModelQueryHandler(service=service, logger=logging.getLogger(__name__))
response = await handler.get_base_models(
SimpleNamespace(query={"limit": "0"}) # pyright: ignore[reportArgumentType]
)
text = response.text
assert text is not None
payload = json.loads(text)
assert payload["success"] is True
assert service.received_limit == 0
@pytest.mark.asyncio
async def test_model_query_handler_rejects_negative_limit_for_base_models():
service = DummyService()
handler = ModelQueryHandler(service=service, logger=logging.getLogger(__name__))
await handler.get_base_models(
SimpleNamespace(query={"limit": "-1"}) # pyright: ignore[reportArgumentType]
)
assert service.received_limit == 20
class DummySearchTagsService:
"""Minimal service stub recording search_tags arguments."""
def __init__(self, result=None):
self.received_query = None
self.received_limit = None
self._result = result or []
async def search_tags(self, query, limit):
self.received_query = query
self.received_limit = limit
return self._result
@pytest.mark.asyncio
async def test_model_query_handler_search_tags_passes_query_and_limit():
service = DummySearchTagsService(result=[{"tag": "anime", "count": 3}])
handler = ModelQueryHandler(service=service, logger=logging.getLogger(__name__))
response = await handler.search_tags(
SimpleNamespace(query={"q": "ani", "limit": "50"}) # pyright: ignore[reportArgumentType]
)
text = response.text
assert text is not None
payload = json.loads(text)
assert payload["success"] is True
assert payload["tags"] == [{"tag": "anime", "count": 3}]
assert service.received_query == "ani"
assert service.received_limit == 50
@pytest.mark.asyncio
async def test_model_query_handler_search_tags_defaults_limit_to_20():
service = DummySearchTagsService()
handler = ModelQueryHandler(service=service, logger=logging.getLogger(__name__))
await handler.search_tags(SimpleNamespace(query={}) # pyright: ignore[reportArgumentType]
)
assert service.received_limit == 20
@pytest.mark.asyncio
async def test_model_query_handler_search_tags_clamps_negative_limit():
service = DummySearchTagsService()
handler = ModelQueryHandler(service=service, logger=logging.getLogger(__name__))
await handler.search_tags(
SimpleNamespace(query={"limit": "-5"}) # pyright: ignore[reportArgumentType]
)
assert service.received_limit == 20
class DummyFolderCache:
def __init__(self, folders):
self.folders = list(folders)
class DummyFolderService:
"""Minimal service stub for the folders/tree endpoints."""
def __init__(self, folders, all_folders):
self.scanner = SimpleNamespace()
cache = DummyFolderCache(folders)
async def get_cached_data(*_, **__):
return cache
async def get_all_folders():
return list(all_folders)
self.scanner.get_cached_data = get_cached_data
self.scanner.get_all_folders = get_all_folders
self.received_include_empty = None
async def get_folder_tree(self, model_root, include_empty=False):
self.received_include_empty = include_empty
return {"tree": "per-root"}
async def get_unified_folder_tree(self, include_empty=False):
self.received_include_empty = include_empty
return {"tree": "unified"}
@pytest.mark.asyncio
async def test_get_folders_defaults_to_models_only_folders():
service = DummyFolderService(["a"], ["a", "empty"])
handler = ModelQueryHandler(service=service, logger=logging.getLogger(__name__))
response = await handler.get_folders(
SimpleNamespace(query={}) # pyright: ignore[reportArgumentType]
)
payload = json.loads(response.text)
assert payload["folders"] == ["a"]
@pytest.mark.asyncio
async def test_get_folders_include_empty_returns_all_folders():
service = DummyFolderService(["a"], ["a", "empty"])
handler = ModelQueryHandler(service=service, logger=logging.getLogger(__name__))
response = await handler.get_folders(
SimpleNamespace(query={"include_empty": "1"}) # pyright: ignore[reportArgumentType]
)
payload = json.loads(response.text)
assert payload["folders"] == ["a", "empty"]
@pytest.mark.asyncio
async def test_get_unified_folder_tree_threads_include_empty():
service = DummyFolderService(["a"], ["a", "empty"])
handler = ModelQueryHandler(service=service, logger=logging.getLogger(__name__))
response = await handler.get_unified_folder_tree(
SimpleNamespace(query={"include_empty": "1"}) # pyright: ignore[reportArgumentType]
)
payload = json.loads(response.text)
assert payload["success"] is True
assert service.received_include_empty is True
@pytest.mark.asyncio
async def test_get_unified_folder_tree_defaults_include_empty_false():
service = DummyFolderService(["a"], ["a", "empty"])
handler = ModelQueryHandler(service=service, logger=logging.getLogger(__name__))
response = await handler.get_unified_folder_tree(
SimpleNamespace(query={}) # pyright: ignore[reportArgumentType]
)
payload = json.loads(response.text)
assert payload["success"] is True
assert service.received_include_empty is False
@pytest.mark.asyncio
async def test_get_folder_tree_threads_include_empty():
service = DummyFolderService(["a"], ["a", "empty"])
handler = ModelQueryHandler(service=service, logger=logging.getLogger(__name__))
response = await handler.get_folder_tree(
SimpleNamespace(query={"model_root": "/models/loras", "include_empty": "true"}) # pyright: ignore[reportArgumentType]
)
payload = json.loads(response.text)
assert payload["success"] is True
assert service.received_include_empty is True
class QueryParams(dict):
"""Minimal stand-in for aiohttp's MultiDict query (supports getall)."""
def getall(self, key, default=None):
value = self.get(key)
if value is None:
return list(default) if default else []
return value if isinstance(value, list) else [value]
class ScopedScanService:
"""Stub service recording the scan scope the handler resolves."""
model_type = "lora"
def __init__(self, roots=None, summary=None, cancelled=False, unavailable=None):
self._roots = list(roots or [])
self._unavailable = list(unavailable or [])
self._summary = summary
self.cancelled = cancelled
self.scanner = SimpleNamespace(is_cancelled=lambda: self.cancelled)
self.received_scope = "not-called"
self.received_rebuild = None
self.root_refreshes = 0
def get_model_roots(self):
return list(self._roots)
def refresh_model_roots(self):
"""Mimic the scanner admitting roots that became readable again."""
self.root_refreshes += 1
return []
def describe_model_roots(self):
return [
{
"path": root,
"label": root.rsplit("/", 1)[-1],
"reachable": True,
"available": True,
"models": 7,
}
for root in self._roots
] + [
{
"path": root,
"label": root.rsplit("/", 1)[-1],
"reachable": False,
"available": False,
"models": 3,
}
for root in self._unavailable
]
async def scan_models(self, force_refresh=False, rebuild_cache=False, scope=None):
self.received_scope = scope
self.received_rebuild = rebuild_cache
return self._summary
SUMMARY = {
"added": 3,
"removed": 0,
"repaired": 1,
"scanned_roots": ["a"],
"skipped_roots": [],
"unavailable_paths": [],
"unavailable_paths_total": 0,
"kept_unreachable": 0,
}
@pytest.mark.asyncio
async def test_scan_models_accepts_roots_param():
service = ScopedScanService(roots=["/mnt/a", "/mnt/b"], summary=SUMMARY)
handler = ModelQueryHandler(service=service, logger=logging.getLogger(__name__))
response = await handler.scan_models(
SimpleNamespace(query=QueryParams({"roots": ["/mnt/a"]}))
)
payload = json.loads(response.text)
assert response.status == 200
assert payload["status"] == "success"
assert payload["added"] == 3
assert payload["scanned_roots"] == ["a"]
assert service.received_scope is not None
assert service.received_scope.roots == ("/mnt/a",)
assert service.received_rebuild is False
# A drive plugged in after startup is admitted before the roots are validated.
assert service.root_refreshes == 1
@pytest.mark.asyncio
async def test_scan_models_without_roots_scans_every_root():
service = ScopedScanService(roots=["/mnt/a"], summary=SUMMARY)
handler = ModelQueryHandler(service=service, logger=logging.getLogger(__name__))
await handler.scan_models(SimpleNamespace(query=QueryParams({})))
assert service.received_scope is None
@pytest.mark.asyncio
async def test_scan_models_rejects_unknown_root():
service = ScopedScanService(roots=["/mnt/a"], summary=SUMMARY)
handler = ModelQueryHandler(service=service, logger=logging.getLogger(__name__))
response = await handler.scan_models(
SimpleNamespace(query=QueryParams({"roots": ["/mnt/nope"]}))
)
payload = json.loads(response.text)
assert response.status == 400
assert payload["roots"] == ["/mnt/nope"]
assert service.received_scope == "not-called"
@pytest.mark.asyncio
async def test_scan_models_accepts_a_root_in_the_hosts_own_spelling():
"""Windows callers may pass a root the way ComfyUI spells it (issue #1146).
``settings.json`` and the ComfyUI logs show ``E:\\loras`` while the live list
holds ``E:/loras``; both name the same directory, so a scoped scan must select
it instead of answering 400 "Unknown model root(s)".
"""
service = ScopedScanService(
roots=["E:/StableDiffusion/cc_models/loras"], summary=SUMMARY
)
handler = ModelQueryHandler(service=service, logger=logging.getLogger(__name__))
for requested in (
"E:\\StableDiffusion\\cc_models\\loras",
"e:/StableDiffusion/cc_models/loras",
):
service.received_scope = "not-called"
response = await handler.scan_models(
SimpleNamespace(query=QueryParams({"roots": [requested]}))
)
assert response.status == 200, requested
assert service.received_scope.roots == (requested,)
@pytest.mark.asyncio
async def test_scan_models_rejects_roots_with_full_rebuild():
service = ScopedScanService(roots=["/mnt/a"], summary=SUMMARY)
handler = ModelQueryHandler(service=service, logger=logging.getLogger(__name__))
response = await handler.scan_models(
SimpleNamespace(
query=QueryParams({"roots": ["/mnt/a"], "full_rebuild": "true"})
)
)
assert response.status == 400
assert service.received_scope == "not-called"
@pytest.mark.asyncio
async def test_get_model_roots_reports_configured_but_unavailable_roots():
"""Wave 6: a root whose directory was missing at startup is still reported."""
service = ScopedScanService(
roots=["/mnt/a"], unavailable=["/mnt/drive-Z/loras"]
)
handler = ModelQueryHandler(service=service, logger=logging.getLogger(__name__))
response = await handler.get_model_roots(SimpleNamespace(query=QueryParams({})))
payload = json.loads(response.text)
assert response.status == 200
assert payload["roots"] == ["/mnt/a"]
details = {entry["path"]: entry for entry in payload["root_details"]}
assert details["/mnt/a"]["available"] is True
assert details["/mnt/drive-Z/loras"]["available"] is False
assert details["/mnt/drive-Z/loras"]["reachable"] is False
assert details["/mnt/drive-Z/loras"]["models"] == 3
# Describing the roots admits anything that became readable again.
assert service.root_refreshes == 1
@pytest.mark.asyncio
async def test_scan_models_accepts_folder_param():
service = ScopedScanService(roots=["/mnt/a", "/mnt/b"], summary=SUMMARY)
handler = ModelQueryHandler(service=service, logger=logging.getLogger(__name__))
response = await handler.scan_models(
SimpleNamespace(query=QueryParams({"folder": "pack\\sub"}))
)
assert response.status == 200
assert service.received_scope is not None
# Backslashes are normalized and no root is pinned: the backend walks this
# relative folder under every root that holds it.
assert service.received_scope.folder == "pack/sub"
assert service.received_scope.roots is None
@pytest.mark.asyncio
async def test_scan_models_rejects_absolute_folder():
service = ScopedScanService(roots=["/mnt/a"], summary=SUMMARY)
handler = ModelQueryHandler(service=service, logger=logging.getLogger(__name__))
for folder in ("/mnt/a/pack", "../pack", "C:/models/pack"):
response = await handler.scan_models(
SimpleNamespace(query=QueryParams({"folder": folder}))
)
assert response.status == 400, folder
assert "error" in json.loads(response.text)
assert service.received_scope == "not-called"
@pytest.mark.asyncio
async def test_scan_models_rejects_folder_with_full_rebuild():
service = ScopedScanService(roots=["/mnt/a"], summary=SUMMARY)
handler = ModelQueryHandler(service=service, logger=logging.getLogger(__name__))
response = await handler.scan_models(
SimpleNamespace(query=QueryParams({"folder": "pack", "full_rebuild": "true"}))
)
assert response.status == 400
assert service.received_scope == "not-called"
@pytest.mark.asyncio
async def test_scan_models_combines_folder_and_roots():
service = ScopedScanService(roots=["/mnt/a", "/mnt/b"], summary=SUMMARY)
handler = ModelQueryHandler(service=service, logger=logging.getLogger(__name__))
response = await handler.scan_models(
SimpleNamespace(query=QueryParams({"folder": "pack", "roots": ["/mnt/b"]}))
)
assert response.status == 200
assert service.received_scope.folder == "pack"
assert service.received_scope.roots == ("/mnt/b",)
@pytest.mark.asyncio
async def test_get_model_roots_reports_details():
service = ScopedScanService(roots=["/mnt/a", "/mnt/b"])
handler = ModelQueryHandler(service=service, logger=logging.getLogger(__name__))
response = await handler.get_model_roots(SimpleNamespace(query=QueryParams({})))
payload = json.loads(response.text)
# `roots` stays a plain list of paths for the existing callers.
assert payload["roots"] == ["/mnt/a", "/mnt/b"]
assert [detail["label"] for detail in payload["root_details"]] == ["a", "b"]
assert all(detail["models"] == 7 for detail in payload["root_details"])