mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-24 06:21:26 -03:00
fix(ui): show empty folders as move and download destinations (#999)
This commit is contained in:
@@ -1030,6 +1030,11 @@ class ModelQueryHandler:
|
|||||||
self._service = service
|
self._service = service
|
||||||
self._logger = logger
|
self._logger = logger
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _parse_include_empty(request: web.Request) -> bool:
|
||||||
|
"""Parse the include_empty query flag (``1``/``true``)."""
|
||||||
|
return request.query.get("include_empty", "").lower() in ("1", "true")
|
||||||
|
|
||||||
async def get_top_tags(self, request: web.Request) -> web.Response:
|
async def get_top_tags(self, request: web.Request) -> web.Response:
|
||||||
try:
|
try:
|
||||||
limit = int(request.query.get("limit", "20"))
|
limit = int(request.query.get("limit", "20"))
|
||||||
@@ -1124,8 +1129,14 @@ class ModelQueryHandler:
|
|||||||
|
|
||||||
async def get_folders(self, request: web.Request) -> web.Response:
|
async def get_folders(self, request: web.Request) -> web.Response:
|
||||||
try:
|
try:
|
||||||
|
include_empty = self._parse_include_empty(request)
|
||||||
|
if include_empty:
|
||||||
|
# Live enumeration includes empty OS-created directories.
|
||||||
|
folders = await self._service.scanner.get_all_folders()
|
||||||
|
else:
|
||||||
cache = await self._service.scanner.get_cached_data()
|
cache = await self._service.scanner.get_cached_data()
|
||||||
return web.json_response({"folders": cache.folders})
|
folders = cache.folders
|
||||||
|
return web.json_response({"folders": folders})
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
self._logger.error("Error getting folders: %s", exc)
|
self._logger.error("Error getting folders: %s", exc)
|
||||||
return web.json_response({"success": False, "error": str(exc)}, status=500)
|
return web.json_response({"success": False, "error": str(exc)}, status=500)
|
||||||
@@ -1150,7 +1161,9 @@ class ModelQueryHandler:
|
|||||||
{"success": False, "error": "model_root parameter is required"},
|
{"success": False, "error": "model_root parameter is required"},
|
||||||
status=400,
|
status=400,
|
||||||
)
|
)
|
||||||
folder_tree = await self._service.get_folder_tree(model_root)
|
folder_tree = await self._service.get_folder_tree(
|
||||||
|
model_root, include_empty=self._parse_include_empty(request)
|
||||||
|
)
|
||||||
return web.json_response({"success": True, "tree": folder_tree})
|
return web.json_response({"success": True, "tree": folder_tree})
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
self._logger.error("Error getting folder tree: %s", exc)
|
self._logger.error("Error getting folder tree: %s", exc)
|
||||||
@@ -1158,7 +1171,9 @@ class ModelQueryHandler:
|
|||||||
|
|
||||||
async def get_unified_folder_tree(self, request: web.Request) -> web.Response:
|
async def get_unified_folder_tree(self, request: web.Request) -> web.Response:
|
||||||
try:
|
try:
|
||||||
unified_tree = await self._service.get_unified_folder_tree()
|
unified_tree = await self._service.get_unified_folder_tree(
|
||||||
|
include_empty=self._parse_include_empty(request)
|
||||||
|
)
|
||||||
return web.json_response({"success": True, "tree": unified_tree})
|
return web.json_response({"success": True, "tree": unified_tree})
|
||||||
except Exception as exc:
|
except Exception as exc:
|
||||||
self._logger.error("Error getting unified folder tree: %s", exc)
|
self._logger.error("Error getting unified folder tree: %s", exc)
|
||||||
|
|||||||
@@ -972,14 +972,25 @@ class BaseModelService(ABC):
|
|||||||
)
|
)
|
||||||
return {k: data[k] for k in fields if k in data}
|
return {k: data[k] for k in fields if k in data}
|
||||||
|
|
||||||
async def get_folder_tree(self, model_root: str) -> Dict[str, Any]:
|
async def _get_tree_folders(self, cache, include_empty: bool) -> List[str]:
|
||||||
|
"""Return the folder list backing folder tree responses.
|
||||||
|
|
||||||
|
With ``include_empty`` the directories are enumerated live from the
|
||||||
|
filesystem (including empty ones) via the scanner; otherwise the
|
||||||
|
models-only ``cache.folders`` list is used unchanged.
|
||||||
|
"""
|
||||||
|
if include_empty:
|
||||||
|
return await self.scanner.get_all_folders()
|
||||||
|
return cache.folders
|
||||||
|
|
||||||
|
async def get_folder_tree(self, model_root: str, include_empty: bool = False) -> Dict[str, Any]:
|
||||||
"""Get hierarchical folder tree for a specific model root"""
|
"""Get hierarchical folder tree for a specific model root"""
|
||||||
cache = await self.scanner.get_cached_data()
|
cache = await self.scanner.get_cached_data()
|
||||||
|
|
||||||
# Build tree structure from folders
|
# Build tree structure from folders
|
||||||
tree = {}
|
tree = {}
|
||||||
|
|
||||||
for folder in cache.folders:
|
for folder in await self._get_tree_folders(cache, include_empty):
|
||||||
# Check if this folder belongs to the specified model root
|
# Check if this folder belongs to the specified model root
|
||||||
folder_belongs_to_root = False
|
folder_belongs_to_root = False
|
||||||
for root in self.scanner.get_model_roots():
|
for root in self.scanner.get_model_roots():
|
||||||
@@ -1001,7 +1012,7 @@ class BaseModelService(ABC):
|
|||||||
|
|
||||||
return tree
|
return tree
|
||||||
|
|
||||||
async def get_unified_folder_tree(self) -> Dict[str, Any]:
|
async def get_unified_folder_tree(self, include_empty: bool = False) -> Dict[str, Any]:
|
||||||
"""Get unified folder tree across all model roots"""
|
"""Get unified folder tree across all model roots"""
|
||||||
cache = await self.scanner.get_cached_data()
|
cache = await self.scanner.get_cached_data()
|
||||||
|
|
||||||
@@ -1011,7 +1022,7 @@ class BaseModelService(ABC):
|
|||||||
# Get all model roots for path normalization
|
# Get all model roots for path normalization
|
||||||
model_roots = self.scanner.get_model_roots()
|
model_roots = self.scanner.get_model_roots()
|
||||||
|
|
||||||
for folder in cache.folders:
|
for folder in await self._get_tree_folders(cache, include_empty):
|
||||||
if not folder: # Skip empty folders
|
if not folder: # Skip empty folders
|
||||||
continue
|
continue
|
||||||
|
|
||||||
|
|||||||
@@ -5,7 +5,7 @@ import asyncio
|
|||||||
import time
|
import time
|
||||||
import shutil
|
import shutil
|
||||||
from dataclasses import dataclass
|
from dataclasses import dataclass
|
||||||
from typing import Any, Awaitable, Callable, Dict, List, Mapping, Optional, Sequence, Set, Type, Union, cast
|
from typing import Any, Awaitable, Callable, Dict, List, Mapping, Optional, Sequence, Set, Tuple, Type, Union, cast
|
||||||
|
|
||||||
from ..utils.models import BaseModelMetadata, autov3_from_civitai_files
|
from ..utils.models import BaseModelMetadata, autov3_from_civitai_files
|
||||||
from ..config import config
|
from ..config import config
|
||||||
@@ -57,6 +57,16 @@ def _is_excluded_dir(name: str) -> bool:
|
|||||||
return name == PENDING_DELETE_DIR_NAME
|
return name == PENDING_DELETE_DIR_NAME
|
||||||
|
|
||||||
|
|
||||||
|
def _is_hidden_relative_path(rel_path: str) -> bool:
|
||||||
|
"""Return True when any segment of a relative path is a hidden directory."""
|
||||||
|
return any(part.startswith(".") for part in rel_path.replace(os.sep, "/").split("/"))
|
||||||
|
|
||||||
|
|
||||||
|
# TTL (seconds) for the get_all_folders() live-walk cache, so rapid repeated
|
||||||
|
# requests (modal open + autocomplete) do not re-walk the model roots.
|
||||||
|
ALL_FOLDERS_CACHE_TTL_SECONDS = 5.0
|
||||||
|
|
||||||
|
|
||||||
def _is_pending_delete_path(path: str) -> bool:
|
def _is_pending_delete_path(path: str) -> bool:
|
||||||
"""Return True when any path component is the pending-delete staging dir."""
|
"""Return True when any path component is the pending-delete staging dir."""
|
||||||
normalized = str(path).replace(os.sep, "/")
|
normalized = str(path).replace(os.sep, "/")
|
||||||
@@ -126,6 +136,8 @@ class ModelScanner:
|
|||||||
self._name_display_mode = self._resolve_name_display_mode()
|
self._name_display_mode = self._resolve_name_display_mode()
|
||||||
self._cancel_requested = False # Flag for cancellation
|
self._cancel_requested = False # Flag for cancellation
|
||||||
self._autov3_backfill_scheduled = False # One-time AutoV3 backfill trigger per process
|
self._autov3_backfill_scheduled = False # One-time AutoV3 backfill trigger per process
|
||||||
|
# Short-lived cache for get_all_folders(): (timestamp, folders) or None
|
||||||
|
self._all_folders_ttl_cache: Optional[Tuple[float, List[str]]] = None
|
||||||
try:
|
try:
|
||||||
loop = asyncio.get_running_loop()
|
loop = asyncio.get_running_loop()
|
||||||
except RuntimeError:
|
except RuntimeError:
|
||||||
@@ -165,6 +177,7 @@ class ModelScanner:
|
|||||||
self._excluded_models = []
|
self._excluded_models = []
|
||||||
self._is_initializing = False
|
self._is_initializing = False
|
||||||
self._name_display_mode = self._resolve_name_display_mode()
|
self._name_display_mode = self._resolve_name_display_mode()
|
||||||
|
self.invalidate_all_folders_cache()
|
||||||
self.bump_cache_version()
|
self.bump_cache_version()
|
||||||
|
|
||||||
try:
|
try:
|
||||||
@@ -1115,6 +1128,56 @@ class ModelScanner:
|
|||||||
"""Get model root directories"""
|
"""Get model root directories"""
|
||||||
raise NotImplementedError("Subclasses must implement get_model_roots")
|
raise NotImplementedError("Subclasses must implement get_model_roots")
|
||||||
|
|
||||||
|
async def get_all_folders(self) -> List[str]:
|
||||||
|
"""Enumerate every directory under the model roots, live from disk.
|
||||||
|
|
||||||
|
Unlike the models-only ``cache.folders``, this includes empty
|
||||||
|
directories, so it stays accurate even when the in-memory cache was
|
||||||
|
hydrated from a persisted snapshot without a filesystem walk. Hidden
|
||||||
|
directories (any segment starting with '.') and the pending-delete
|
||||||
|
staging dir are excluded. The result is unioned with the model-derived
|
||||||
|
folders so it is always a superset of ``cache.folders``, and cached
|
||||||
|
for ``ALL_FOLDERS_CACHE_TTL_SECONDS`` to avoid repeated walks.
|
||||||
|
"""
|
||||||
|
now = time.monotonic()
|
||||||
|
if self._all_folders_ttl_cache is not None:
|
||||||
|
cached_at, cached_folders = self._all_folders_ttl_cache
|
||||||
|
if now - cached_at < ALL_FOLDERS_CACHE_TTL_SECONDS:
|
||||||
|
return cached_folders
|
||||||
|
|
||||||
|
discovered: Set[str] = set()
|
||||||
|
visited_real_paths: Set[str] = set()
|
||||||
|
|
||||||
|
for root_path in self.get_model_roots():
|
||||||
|
if not os.path.exists(root_path):
|
||||||
|
continue
|
||||||
|
|
||||||
|
for root, dirnames, _files in os.walk(root_path, followlinks=True):
|
||||||
|
dirnames[:] = [d for d in dirnames if not _is_excluded_dir(d)]
|
||||||
|
# realpath is used only for symlink dedup, never for the
|
||||||
|
# recorded path (business paths stay unresolved).
|
||||||
|
real_root = os.path.realpath(root)
|
||||||
|
if real_root in visited_real_paths:
|
||||||
|
continue
|
||||||
|
visited_real_paths.add(real_root)
|
||||||
|
|
||||||
|
rel_dir = os.path.relpath(os.path.abspath(root), os.path.abspath(root_path))
|
||||||
|
rel_dir = rel_dir.replace(os.path.sep, "/")
|
||||||
|
if rel_dir != "." and not _is_hidden_relative_path(rel_dir):
|
||||||
|
discovered.add(rel_dir)
|
||||||
|
|
||||||
|
folders = set(discovered)
|
||||||
|
if self._cache is not None:
|
||||||
|
folders |= {item.get('folder', '') for item in self._cache.raw_data}
|
||||||
|
|
||||||
|
result = sorted(folders, key=lambda x: x.lower())
|
||||||
|
self._all_folders_ttl_cache = (now, result)
|
||||||
|
return result
|
||||||
|
|
||||||
|
def invalidate_all_folders_cache(self) -> None:
|
||||||
|
"""Drop the cached get_all_folders() result (e.g. after a move)."""
|
||||||
|
self._all_folders_ttl_cache = None
|
||||||
|
|
||||||
async def _create_default_metadata(self, file_path: str) -> Optional[BaseModelMetadata]:
|
async def _create_default_metadata(self, file_path: str) -> Optional[BaseModelMetadata]:
|
||||||
"""Get model file info and metadata (extensible for different model types)"""
|
"""Get model file info and metadata (extensible for different model types)"""
|
||||||
return await MetadataManager.create_default_metadata(file_path, self.model_class)
|
return await MetadataManager.create_default_metadata(file_path, self.model_class)
|
||||||
@@ -1773,6 +1836,10 @@ class ModelScanner:
|
|||||||
|
|
||||||
await cache.resort()
|
await cache.resort()
|
||||||
|
|
||||||
|
# A move may have created new directories; drop the cached live-walk
|
||||||
|
# result so the next include_empty request sees them.
|
||||||
|
self.invalidate_all_folders_cache()
|
||||||
|
|
||||||
if cache_modified:
|
if cache_modified:
|
||||||
await self._persist_current_cache()
|
await self._persist_current_cache()
|
||||||
self.bump_cache_version()
|
self.bump_cache_version()
|
||||||
|
|||||||
@@ -1206,9 +1206,13 @@ export class BaseModelApiClient {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
async fetchUnifiedFolderTree() {
|
async fetchUnifiedFolderTree(options = {}) {
|
||||||
try {
|
try {
|
||||||
const response = await fetch(this.apiConfig.endpoints.unifiedFolderTree);
|
const { includeEmpty = false } = options;
|
||||||
|
const url = includeEmpty
|
||||||
|
? `${this.apiConfig.endpoints.unifiedFolderTree}?include_empty=1`
|
||||||
|
: this.apiConfig.endpoints.unifiedFolderTree;
|
||||||
|
const response = await fetch(url);
|
||||||
if (!response.ok) {
|
if (!response.ok) {
|
||||||
throw new Error(`Failed to fetch unified folder tree`);
|
throw new Error(`Failed to fetch unified folder tree`);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -2134,8 +2134,9 @@ export class DownloadManager {
|
|||||||
|
|
||||||
async initializeFolderTree() {
|
async initializeFolderTree() {
|
||||||
try {
|
try {
|
||||||
// Fetch unified folder tree
|
// Fetch unified folder tree, including empty directories so they
|
||||||
const treeData = await this.apiClient.fetchUnifiedFolderTree();
|
// can be selected as download destinations
|
||||||
|
const treeData = await this.apiClient.fetchUnifiedFolderTree({ includeEmpty: true });
|
||||||
|
|
||||||
if (treeData.success) {
|
if (treeData.success) {
|
||||||
// Load tree data into folder tree manager
|
// Load tree data into folder tree manager
|
||||||
|
|||||||
@@ -200,8 +200,9 @@ class MoveManager {
|
|||||||
async initializeFolderTree() {
|
async initializeFolderTree() {
|
||||||
try {
|
try {
|
||||||
const apiClient = this._getApiClient();
|
const apiClient = this._getApiClient();
|
||||||
// Fetch unified folder tree
|
// Fetch unified folder tree, including empty directories so they
|
||||||
const treeData = await apiClient.fetchUnifiedFolderTree();
|
// can be selected as move targets
|
||||||
|
const treeData = await apiClient.fetchUnifiedFolderTree({ includeEmpty: true });
|
||||||
|
|
||||||
if (treeData.success) {
|
if (treeData.success) {
|
||||||
// Load tree data into folder tree manager
|
// Load tree data into folder tree manager
|
||||||
|
|||||||
@@ -0,0 +1,97 @@
|
|||||||
|
import { describe, it, afterEach, expect, vi } from 'vitest';
|
||||||
|
|
||||||
|
const {
|
||||||
|
BASE_MODEL_API_MODULE,
|
||||||
|
STATE_MODULE,
|
||||||
|
UI_HELPERS_MODULE,
|
||||||
|
I18N_MODULE,
|
||||||
|
STORAGE_MODULE,
|
||||||
|
API_CONFIG_MODULE,
|
||||||
|
API_FACTORY_MODULE,
|
||||||
|
SIDEBAR_MANAGER_MODULE,
|
||||||
|
} = vi.hoisted(() => ({
|
||||||
|
BASE_MODEL_API_MODULE: new URL('../../../static/js/api/baseModelApi.js', import.meta.url).pathname,
|
||||||
|
STATE_MODULE: new URL('../../../static/js/state/index.js', import.meta.url).pathname,
|
||||||
|
UI_HELPERS_MODULE: new URL('../../../static/js/utils/uiHelpers.js', import.meta.url).pathname,
|
||||||
|
I18N_MODULE: new URL('../../../static/js/utils/i18nHelpers.js', import.meta.url).pathname,
|
||||||
|
STORAGE_MODULE: new URL('../../../static/js/utils/storageHelpers.js', import.meta.url).pathname,
|
||||||
|
API_CONFIG_MODULE: new URL('../../../static/js/api/apiConfig.js', import.meta.url).pathname,
|
||||||
|
API_FACTORY_MODULE: new URL('../../../static/js/api/modelApiFactory.js', import.meta.url).pathname,
|
||||||
|
SIDEBAR_MANAGER_MODULE: new URL('../../../static/js/components/SidebarManager.js', import.meta.url).pathname,
|
||||||
|
}));
|
||||||
|
|
||||||
|
vi.mock(STATE_MODULE, () => ({
|
||||||
|
state: {},
|
||||||
|
getCurrentPageState: vi.fn(() => ({})),
|
||||||
|
}));
|
||||||
|
|
||||||
|
vi.mock(UI_HELPERS_MODULE, () => ({
|
||||||
|
showToast: vi.fn(),
|
||||||
|
}));
|
||||||
|
|
||||||
|
vi.mock(I18N_MODULE, () => ({
|
||||||
|
translate: vi.fn((key) => key),
|
||||||
|
}));
|
||||||
|
|
||||||
|
vi.mock(STORAGE_MODULE, () => ({
|
||||||
|
getStorageItem: vi.fn(),
|
||||||
|
getSessionItem: vi.fn(),
|
||||||
|
removeSessionItem: vi.fn(),
|
||||||
|
saveMapToStorage: vi.fn(),
|
||||||
|
}));
|
||||||
|
|
||||||
|
vi.mock(API_CONFIG_MODULE, () => ({
|
||||||
|
getCompleteApiConfig: vi.fn(() => ({
|
||||||
|
endpoints: { unifiedFolderTree: '/api/lm/loras/unified-folder-tree' },
|
||||||
|
config: { displayName: 'LoRA', singularName: 'LoRA' },
|
||||||
|
})),
|
||||||
|
getCurrentModelType: vi.fn(() => 'loras'),
|
||||||
|
isValidModelType: vi.fn(() => true),
|
||||||
|
DOWNLOAD_ENDPOINTS: {},
|
||||||
|
HF_ENDPOINTS: {},
|
||||||
|
WS_ENDPOINTS: {},
|
||||||
|
}));
|
||||||
|
|
||||||
|
vi.mock(API_FACTORY_MODULE, () => ({
|
||||||
|
resetAndReload: vi.fn(),
|
||||||
|
}));
|
||||||
|
|
||||||
|
vi.mock(SIDEBAR_MANAGER_MODULE, () => ({
|
||||||
|
sidebarManager: { refresh: vi.fn() },
|
||||||
|
}));
|
||||||
|
|
||||||
|
describe('BaseModelApiClient.fetchUnifiedFolderTree', () => {
|
||||||
|
afterEach(() => {
|
||||||
|
delete global.fetch;
|
||||||
|
});
|
||||||
|
|
||||||
|
async function createClient() {
|
||||||
|
const { BaseModelApiClient } = await import(BASE_MODEL_API_MODULE);
|
||||||
|
class TestClient extends BaseModelApiClient {}
|
||||||
|
return new TestClient('loras');
|
||||||
|
}
|
||||||
|
|
||||||
|
it('requests the plain endpoint by default', async () => {
|
||||||
|
global.fetch = vi.fn().mockResolvedValue({
|
||||||
|
ok: true,
|
||||||
|
json: async () => ({ success: true, tree: {} }),
|
||||||
|
});
|
||||||
|
|
||||||
|
const client = await createClient();
|
||||||
|
await client.fetchUnifiedFolderTree();
|
||||||
|
|
||||||
|
expect(global.fetch).toHaveBeenCalledWith('/api/lm/loras/unified-folder-tree');
|
||||||
|
});
|
||||||
|
|
||||||
|
it('appends include_empty=1 when requested', async () => {
|
||||||
|
global.fetch = vi.fn().mockResolvedValue({
|
||||||
|
ok: true,
|
||||||
|
json: async () => ({ success: true, tree: {} }),
|
||||||
|
});
|
||||||
|
|
||||||
|
const client = await createClient();
|
||||||
|
await client.fetchUnifiedFolderTree({ includeEmpty: true });
|
||||||
|
|
||||||
|
expect(global.fetch).toHaveBeenCalledWith('/api/lm/loras/unified-folder-tree?include_empty=1');
|
||||||
|
});
|
||||||
|
});
|
||||||
@@ -92,6 +92,12 @@ describe('MoveManager', () => {
|
|||||||
expect(moveManager.folderTreeManager.getSelectedPath()).toBe('');
|
expect(moveManager.folderTreeManager.getSelectedPath()).toBe('');
|
||||||
});
|
});
|
||||||
|
|
||||||
|
it('should fetch the folder tree including empty directories', async () => {
|
||||||
|
await moveManager.initializeFolderTree();
|
||||||
|
|
||||||
|
expect(mockApiClient.fetchUnifiedFolderTree).toHaveBeenCalledWith({ includeEmpty: true });
|
||||||
|
});
|
||||||
|
|
||||||
it('should ignore manual folder selection when useDefaultPath is true', async () => {
|
it('should ignore manual folder selection when useDefaultPath is true', async () => {
|
||||||
// Setup state
|
// Setup state
|
||||||
moveManager.useDefaultPath = true;
|
moveManager.useDefaultPath = true;
|
||||||
|
|||||||
@@ -0,0 +1,134 @@
|
|||||||
|
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
|
||||||
|
|
||||||
|
const {
|
||||||
|
DOWNLOAD_MANAGER_MODULE,
|
||||||
|
MODAL_MANAGER_MODULE,
|
||||||
|
UI_HELPERS_MODULE,
|
||||||
|
STATE_MODULE,
|
||||||
|
LOADING_MANAGER_MODULE,
|
||||||
|
API_FACTORY_MODULE,
|
||||||
|
STORAGE_HELPERS_MODULE,
|
||||||
|
FOLDER_TREE_MANAGER_MODULE,
|
||||||
|
I18N_HELPERS_MODULE,
|
||||||
|
SUMMARY_MODULE,
|
||||||
|
mockApiClient,
|
||||||
|
mockLoadingManager,
|
||||||
|
mockFolderTreeManager,
|
||||||
|
showToastMock,
|
||||||
|
} = vi.hoisted(() => {
|
||||||
|
const mockApiClient = {
|
||||||
|
modelType: 'loras',
|
||||||
|
apiConfig: {
|
||||||
|
config: {
|
||||||
|
displayName: 'LoRA',
|
||||||
|
singularName: 'lora',
|
||||||
|
},
|
||||||
|
},
|
||||||
|
fetchModelRoots: vi.fn(async () => ({ roots: ['/models/loras'] })),
|
||||||
|
fetchUnifiedFolderTree: vi.fn(async () => ({ success: true, tree: {} })),
|
||||||
|
};
|
||||||
|
|
||||||
|
const mockLoadingManager = {
|
||||||
|
showSimpleLoading: vi.fn(),
|
||||||
|
setStatus: vi.fn(),
|
||||||
|
hide: vi.fn(),
|
||||||
|
restoreProgressBar: vi.fn(),
|
||||||
|
showDownloadProgress: vi.fn(() => vi.fn()),
|
||||||
|
showCancelButton: vi.fn(),
|
||||||
|
};
|
||||||
|
|
||||||
|
const mockFolderTreeManager = {
|
||||||
|
clearSelection: vi.fn(),
|
||||||
|
init: vi.fn(),
|
||||||
|
loadTree: vi.fn(async () => {}),
|
||||||
|
getSelectedPath: vi.fn(() => ''),
|
||||||
|
};
|
||||||
|
|
||||||
|
return {
|
||||||
|
DOWNLOAD_MANAGER_MODULE: new URL('../../../static/js/managers/DownloadManager.js', import.meta.url).pathname,
|
||||||
|
MODAL_MANAGER_MODULE: new URL('../../../static/js/managers/ModalManager.js', import.meta.url).pathname,
|
||||||
|
UI_HELPERS_MODULE: new URL('../../../static/js/utils/uiHelpers.js', import.meta.url).pathname,
|
||||||
|
STATE_MODULE: new URL('../../../static/js/state/index.js', import.meta.url).pathname,
|
||||||
|
LOADING_MANAGER_MODULE: new URL('../../../static/js/managers/LoadingManager.js', import.meta.url).pathname,
|
||||||
|
API_FACTORY_MODULE: new URL('../../../static/js/api/modelApiFactory.js', import.meta.url).pathname,
|
||||||
|
STORAGE_HELPERS_MODULE: new URL('../../../static/js/utils/storageHelpers.js', import.meta.url).pathname,
|
||||||
|
FOLDER_TREE_MANAGER_MODULE: new URL('../../../static/js/components/FolderTreeManager.js', import.meta.url).pathname,
|
||||||
|
I18N_HELPERS_MODULE: new URL('../../../static/js/utils/i18nHelpers.js', import.meta.url).pathname,
|
||||||
|
SUMMARY_MODULE: new URL('../../../static/js/components/DownloadBatchSummaryModal.js', import.meta.url).pathname,
|
||||||
|
mockApiClient,
|
||||||
|
mockLoadingManager,
|
||||||
|
mockFolderTreeManager,
|
||||||
|
showToastMock: vi.fn(),
|
||||||
|
};
|
||||||
|
});
|
||||||
|
|
||||||
|
vi.mock(MODAL_MANAGER_MODULE, () => ({
|
||||||
|
modalManager: {
|
||||||
|
showModal: vi.fn(),
|
||||||
|
closeModal: vi.fn(),
|
||||||
|
},
|
||||||
|
}));
|
||||||
|
|
||||||
|
vi.mock(UI_HELPERS_MODULE, () => ({
|
||||||
|
showToast: showToastMock,
|
||||||
|
setupAutoNewlineOnPaste: vi.fn(),
|
||||||
|
}));
|
||||||
|
|
||||||
|
vi.mock(STATE_MODULE, () => ({
|
||||||
|
state: {
|
||||||
|
global: {
|
||||||
|
settings: {},
|
||||||
|
},
|
||||||
|
loadingManager: mockLoadingManager,
|
||||||
|
},
|
||||||
|
}));
|
||||||
|
|
||||||
|
vi.mock(LOADING_MANAGER_MODULE, () => ({
|
||||||
|
LoadingManager: vi.fn(() => mockLoadingManager),
|
||||||
|
}));
|
||||||
|
|
||||||
|
vi.mock(API_FACTORY_MODULE, () => ({
|
||||||
|
getModelApiClient: vi.fn(() => mockApiClient),
|
||||||
|
resetAndReload: vi.fn(),
|
||||||
|
}));
|
||||||
|
|
||||||
|
vi.mock(STORAGE_HELPERS_MODULE, () => ({
|
||||||
|
getStorageItem: vi.fn((_key, defaultValue) => defaultValue),
|
||||||
|
setStorageItem: vi.fn(),
|
||||||
|
}));
|
||||||
|
|
||||||
|
vi.mock(FOLDER_TREE_MANAGER_MODULE, () => ({
|
||||||
|
FolderTreeManager: vi.fn(() => mockFolderTreeManager),
|
||||||
|
}));
|
||||||
|
|
||||||
|
vi.mock(I18N_HELPERS_MODULE, () => ({
|
||||||
|
translate: vi.fn((_, __, fallback) => fallback ?? ''),
|
||||||
|
}));
|
||||||
|
|
||||||
|
vi.mock(SUMMARY_MODULE, () => ({
|
||||||
|
showDownloadBatchSummary: vi.fn(),
|
||||||
|
}));
|
||||||
|
|
||||||
|
describe('DownloadManager folder tree', () => {
|
||||||
|
let DownloadManager;
|
||||||
|
|
||||||
|
beforeEach(async () => {
|
||||||
|
vi.resetModules();
|
||||||
|
vi.clearAllMocks();
|
||||||
|
({ DownloadManager } = await import(DOWNLOAD_MANAGER_MODULE));
|
||||||
|
});
|
||||||
|
|
||||||
|
afterEach(() => {
|
||||||
|
document.body.innerHTML = '';
|
||||||
|
});
|
||||||
|
|
||||||
|
it('should fetch the folder tree including empty directories', async () => {
|
||||||
|
const manager = new DownloadManager();
|
||||||
|
manager.apiClient = mockApiClient;
|
||||||
|
|
||||||
|
await manager.initializeFolderTree();
|
||||||
|
|
||||||
|
expect(mockApiClient.fetchUnifiedFolderTree).toHaveBeenCalledWith({ includeEmpty: true });
|
||||||
|
expect(mockFolderTreeManager.loadTree).toHaveBeenCalledWith({});
|
||||||
|
});
|
||||||
|
});
|
||||||
@@ -97,3 +97,102 @@ async def test_model_query_handler_search_tags_clamps_negative_limit():
|
|||||||
)
|
)
|
||||||
|
|
||||||
assert service.received_limit == 20
|
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
|
||||||
|
|||||||
@@ -1394,3 +1394,71 @@ class TestApplyHashFilters:
|
|||||||
result = await service._apply_hash_filters(data, {})
|
result = await service._apply_hash_filters(data, {})
|
||||||
|
|
||||||
assert result == data
|
assert result == data
|
||||||
|
|
||||||
|
|
||||||
|
class FolderTreeCache:
|
||||||
|
def __init__(self, folders):
|
||||||
|
self.folders = list(folders)
|
||||||
|
|
||||||
|
|
||||||
|
class FolderTreeScanner:
|
||||||
|
def __init__(self, cache, all_folders=None):
|
||||||
|
self._cache = cache
|
||||||
|
self._all_folders = list(all_folders) if all_folders is not None else None
|
||||||
|
|
||||||
|
async def get_cached_data(self, *_, **__):
|
||||||
|
return self._cache
|
||||||
|
|
||||||
|
async def get_all_folders(self):
|
||||||
|
if self._all_folders is None:
|
||||||
|
raise AssertionError("get_all_folders should not be called")
|
||||||
|
return list(self._all_folders)
|
||||||
|
|
||||||
|
def get_model_roots(self):
|
||||||
|
return ["/models/loras"]
|
||||||
|
|
||||||
|
|
||||||
|
def _make_folder_tree_service(folders, all_folders=None):
|
||||||
|
cache = FolderTreeCache(folders)
|
||||||
|
scanner = FolderTreeScanner(cache, all_folders)
|
||||||
|
settings = StubSettings({})
|
||||||
|
return DummyService(
|
||||||
|
model_type="stub",
|
||||||
|
scanner=scanner,
|
||||||
|
metadata_class=BaseModelMetadata,
|
||||||
|
cache_repository=ModelCacheRepository(scanner),
|
||||||
|
filter_set=ModelFilterSet(settings),
|
||||||
|
search_strategy=SearchStrategy(),
|
||||||
|
settings_provider=settings,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_get_unified_folder_tree_defaults_to_models_only_folders():
|
||||||
|
service = _make_folder_tree_service(["a/b"])
|
||||||
|
|
||||||
|
tree = await service.get_unified_folder_tree()
|
||||||
|
|
||||||
|
assert tree == {"a": {"b": {}}}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_get_unified_folder_tree_include_empty_uses_live_enumeration():
|
||||||
|
service = _make_folder_tree_service(["a/b"], ["a/b", "empty", "empty/dir"])
|
||||||
|
|
||||||
|
tree = await service.get_unified_folder_tree(include_empty=True)
|
||||||
|
|
||||||
|
assert tree == {"a": {"b": {}}, "empty": {"dir": {}}}
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_get_folder_tree_include_empty_uses_live_enumeration():
|
||||||
|
service = _make_folder_tree_service(["a/b"], ["a/b", "empty"])
|
||||||
|
|
||||||
|
default_tree = await service.get_folder_tree("/models/loras")
|
||||||
|
include_empty_tree = await service.get_folder_tree(
|
||||||
|
"/models/loras", include_empty=True
|
||||||
|
)
|
||||||
|
|
||||||
|
assert default_tree == {"a": {"b": {}}}
|
||||||
|
assert include_empty_tree == {"a": {"b": {}}, "empty": {}}
|
||||||
|
|||||||
@@ -1210,3 +1210,111 @@ async def test_bulk_delete_cancelled_after_one_staged_batch_present(
|
|||||||
assert not first.exists()
|
assert not first.exists()
|
||||||
# The second file was never touched.
|
# The second file was never touched.
|
||||||
assert second.exists()
|
assert second.exists()
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_get_all_folders_enumerates_empty_directories_live(tmp_path: Path):
|
||||||
|
_create_files(tmp_path)
|
||||||
|
(tmp_path / "empty").mkdir()
|
||||||
|
(tmp_path / "empty" / "nested_empty").mkdir()
|
||||||
|
(tmp_path / ".hidden").mkdir()
|
||||||
|
(tmp_path / "visible" / ".hidden_child").mkdir(parents=True)
|
||||||
|
(tmp_path / PENDING_DELETE_DIR_NAME).mkdir()
|
||||||
|
|
||||||
|
scanner = DummyScanner(tmp_path)
|
||||||
|
await scanner._initialize_cache()
|
||||||
|
cache = await scanner.get_cached_data()
|
||||||
|
|
||||||
|
all_folders = await scanner.get_all_folders()
|
||||||
|
|
||||||
|
# cache.folders stays models-only
|
||||||
|
assert sorted(cache.folders) == ["", "nested"]
|
||||||
|
|
||||||
|
# Live enumeration includes empty directories and stays a superset
|
||||||
|
assert set(cache.folders) <= set(all_folders)
|
||||||
|
assert "empty" in all_folders
|
||||||
|
assert "empty/nested_empty" in all_folders
|
||||||
|
assert "visible" in all_folders
|
||||||
|
|
||||||
|
# Hidden directories (any segment starting with '.') are excluded
|
||||||
|
assert not any(
|
||||||
|
segment.startswith(".")
|
||||||
|
for folder in all_folders
|
||||||
|
for segment in folder.split("/")
|
||||||
|
)
|
||||||
|
# The pending-delete staging dir is excluded
|
||||||
|
assert PENDING_DELETE_DIR_NAME not in all_folders
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_get_all_folders_uses_ttl_cache(tmp_path: Path, monkeypatch):
|
||||||
|
_create_files(tmp_path)
|
||||||
|
scanner = DummyScanner(tmp_path)
|
||||||
|
await scanner._initialize_cache()
|
||||||
|
|
||||||
|
walk_calls = {"n": 0}
|
||||||
|
real_walk = os.walk
|
||||||
|
|
||||||
|
def counting_walk(*args, **kwargs):
|
||||||
|
walk_calls["n"] += 1
|
||||||
|
return real_walk(*args, **kwargs)
|
||||||
|
|
||||||
|
monkeypatch.setattr(model_scanner.os, "walk", counting_walk)
|
||||||
|
|
||||||
|
first = await scanner.get_all_folders()
|
||||||
|
assert walk_calls["n"] == 1
|
||||||
|
|
||||||
|
# Second call within the TTL reuses the cached result without re-walking
|
||||||
|
second = await scanner.get_all_folders()
|
||||||
|
assert walk_calls["n"] == 1
|
||||||
|
assert second == first
|
||||||
|
|
||||||
|
# After the TTL expires the roots are walked again
|
||||||
|
real_monotonic = time.monotonic
|
||||||
|
monkeypatch.setattr(
|
||||||
|
model_scanner.time,
|
||||||
|
"monotonic",
|
||||||
|
lambda: real_monotonic() + model_scanner.ALL_FOLDERS_CACHE_TTL_SECONDS + 1,
|
||||||
|
)
|
||||||
|
third = await scanner.get_all_folders()
|
||||||
|
assert walk_calls["n"] == 2
|
||||||
|
assert third == first
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.mark.asyncio
|
||||||
|
async def test_get_all_folders_invalidated_after_move(tmp_path: Path):
|
||||||
|
first, _, _ = _create_files(tmp_path)
|
||||||
|
scanner = DummyScanner(tmp_path)
|
||||||
|
|
||||||
|
await scanner._initialize_cache()
|
||||||
|
|
||||||
|
cached = await scanner.get_all_folders()
|
||||||
|
assert scanner._all_folders_ttl_cache is not None
|
||||||
|
assert "new/deep" not in cached
|
||||||
|
|
||||||
|
# Simulate a move: target directories exist on disk (created by
|
||||||
|
# os.makedirs in move_model) and the cache entry is relocated.
|
||||||
|
(tmp_path / "new" / "deep").mkdir(parents=True)
|
||||||
|
original = _normalize_path(first)
|
||||||
|
new_path = _normalize_path(tmp_path / "new" / "deep" / "one.txt")
|
||||||
|
moved_metadata = {
|
||||||
|
"file_path": new_path,
|
||||||
|
"file_name": "one",
|
||||||
|
"model_name": "one",
|
||||||
|
"sha256": "hash-one",
|
||||||
|
"tags": ["alpha"],
|
||||||
|
"size": 1,
|
||||||
|
"modified": 1.0,
|
||||||
|
}
|
||||||
|
|
||||||
|
await scanner.update_single_model_cache(original, new_path, moved_metadata)
|
||||||
|
|
||||||
|
# The TTL cache was invalidated by the move
|
||||||
|
assert scanner._all_folders_ttl_cache is None
|
||||||
|
|
||||||
|
all_folders = await scanner.get_all_folders()
|
||||||
|
cache = await scanner.get_cached_data()
|
||||||
|
assert sorted(cache.folders) == ["nested", "new/deep"]
|
||||||
|
assert "new" in all_folders
|
||||||
|
assert "new/deep" in all_folders
|
||||||
|
assert set(cache.folders) <= set(all_folders)
|
||||||
|
|||||||
Reference in New Issue
Block a user