diff --git a/py/routes/handlers/model_handlers.py b/py/routes/handlers/model_handlers.py index 3fc2c0b6..22c00ee6 100644 --- a/py/routes/handlers/model_handlers.py +++ b/py/routes/handlers/model_handlers.py @@ -1030,6 +1030,11 @@ class ModelQueryHandler: self._service = service 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: try: limit = int(request.query.get("limit", "20")) @@ -1124,8 +1129,14 @@ class ModelQueryHandler: async def get_folders(self, request: web.Request) -> web.Response: try: - cache = await self._service.scanner.get_cached_data() - return web.json_response({"folders": cache.folders}) + 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() + folders = cache.folders + return web.json_response({"folders": folders}) except Exception as exc: self._logger.error("Error getting folders: %s", exc) 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"}, 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}) except Exception as 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: 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}) except Exception as exc: self._logger.error("Error getting unified folder tree: %s", exc) diff --git a/py/services/base_model_service.py b/py/services/base_model_service.py index 2a250a78..45a762ba 100644 --- a/py/services/base_model_service.py +++ b/py/services/base_model_service.py @@ -972,14 +972,25 @@ class BaseModelService(ABC): ) 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""" cache = await self.scanner.get_cached_data() # Build tree structure from folders 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 folder_belongs_to_root = False for root in self.scanner.get_model_roots(): @@ -1001,7 +1012,7 @@ class BaseModelService(ABC): 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""" cache = await self.scanner.get_cached_data() @@ -1011,7 +1022,7 @@ class BaseModelService(ABC): # Get all model roots for path normalization 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 continue diff --git a/py/services/model_scanner.py b/py/services/model_scanner.py index 05a20277..3baddd7e 100644 --- a/py/services/model_scanner.py +++ b/py/services/model_scanner.py @@ -5,7 +5,7 @@ import asyncio import time import shutil 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 ..config import config @@ -57,6 +57,16 @@ def _is_excluded_dir(name: str) -> bool: 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: """Return True when any path component is the pending-delete staging dir.""" normalized = str(path).replace(os.sep, "/") @@ -126,6 +136,8 @@ class ModelScanner: self._name_display_mode = self._resolve_name_display_mode() self._cancel_requested = False # Flag for cancellation 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: loop = asyncio.get_running_loop() except RuntimeError: @@ -165,6 +177,7 @@ class ModelScanner: self._excluded_models = [] self._is_initializing = False self._name_display_mode = self._resolve_name_display_mode() + self.invalidate_all_folders_cache() self.bump_cache_version() try: @@ -897,12 +910,12 @@ class ModelScanner: new_files = [] visited_real_paths = set() discovered_real_files = set() - + # Scan all model roots for root_path in self.get_model_roots(): if not os.path.exists(root_path): continue - + # Recursively scan directory for root, dirnames, files in os.walk(root_path, followlinks=True): dirnames[:] = [d for d in dirnames if not _is_excluded_dir(d)] @@ -910,7 +923,7 @@ class ModelScanner: if real_root in visited_real_paths: continue visited_real_paths.add(real_root) - + for file in files: ext = os.path.splitext(file)[1].lower() if ext in self.file_extensions: @@ -955,7 +968,7 @@ class ModelScanner: if self.is_cancelled(): logger.info(f"{self.model_type.capitalize()} Scanner: Reconcile scan cancelled") return - + # Process new files in batches total_added = 0 if new_files: @@ -1114,6 +1127,56 @@ class ModelScanner: def get_model_roots(self) -> List[str]: """Get model root directories""" 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]: """Get model file info and metadata (extensible for different model types)""" @@ -1773,6 +1836,10 @@ class ModelScanner: 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: await self._persist_current_cache() self.bump_cache_version() diff --git a/static/js/api/baseModelApi.js b/static/js/api/baseModelApi.js index 3a5bc3ea..f2624e06 100644 --- a/static/js/api/baseModelApi.js +++ b/static/js/api/baseModelApi.js @@ -1206,9 +1206,13 @@ export class BaseModelApiClient { } } - async fetchUnifiedFolderTree() { + async fetchUnifiedFolderTree(options = {}) { 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) { throw new Error(`Failed to fetch unified folder tree`); } diff --git a/static/js/managers/DownloadManager.js b/static/js/managers/DownloadManager.js index 04f776f0..4de4dac9 100644 --- a/static/js/managers/DownloadManager.js +++ b/static/js/managers/DownloadManager.js @@ -2134,8 +2134,9 @@ export class DownloadManager { async initializeFolderTree() { try { - // Fetch unified folder tree - const treeData = await this.apiClient.fetchUnifiedFolderTree(); + // Fetch unified folder tree, including empty directories so they + // can be selected as download destinations + const treeData = await this.apiClient.fetchUnifiedFolderTree({ includeEmpty: true }); if (treeData.success) { // Load tree data into folder tree manager diff --git a/static/js/managers/MoveManager.js b/static/js/managers/MoveManager.js index 400038e5..f6d55f34 100644 --- a/static/js/managers/MoveManager.js +++ b/static/js/managers/MoveManager.js @@ -200,8 +200,9 @@ class MoveManager { async initializeFolderTree() { try { const apiClient = this._getApiClient(); - // Fetch unified folder tree - const treeData = await apiClient.fetchUnifiedFolderTree(); + // Fetch unified folder tree, including empty directories so they + // can be selected as move targets + const treeData = await apiClient.fetchUnifiedFolderTree({ includeEmpty: true }); if (treeData.success) { // Load tree data into folder tree manager diff --git a/tests/frontend/api/baseModelApi.folderTree.test.js b/tests/frontend/api/baseModelApi.folderTree.test.js new file mode 100644 index 00000000..d772ec9f --- /dev/null +++ b/tests/frontend/api/baseModelApi.folderTree.test.js @@ -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'); + }); +}); diff --git a/tests/frontend/managers/MoveManager.test.js b/tests/frontend/managers/MoveManager.test.js index ab829493..1fd78b66 100644 --- a/tests/frontend/managers/MoveManager.test.js +++ b/tests/frontend/managers/MoveManager.test.js @@ -92,6 +92,12 @@ describe('MoveManager', () => { 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 () => { // Setup state moveManager.useDefaultPath = true; diff --git a/tests/frontend/managers/downloadManager.folderTree.test.js b/tests/frontend/managers/downloadManager.folderTree.test.js new file mode 100644 index 00000000..1df0959b --- /dev/null +++ b/tests/frontend/managers/downloadManager.folderTree.test.js @@ -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({}); + }); +}); diff --git a/tests/routes/test_model_query_handler.py b/tests/routes/test_model_query_handler.py index ea88b90f..d3cd27a3 100644 --- a/tests/routes/test_model_query_handler.py +++ b/tests/routes/test_model_query_handler.py @@ -97,3 +97,102 @@ async def test_model_query_handler_search_tags_clamps_negative_limit(): ) 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 diff --git a/tests/services/test_base_model_service.py b/tests/services/test_base_model_service.py index cb6dc1c2..b4ee4ba1 100644 --- a/tests/services/test_base_model_service.py +++ b/tests/services/test_base_model_service.py @@ -1394,3 +1394,71 @@ class TestApplyHashFilters: result = await service._apply_hash_filters(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": {}} diff --git a/tests/services/test_model_scanner.py b/tests/services/test_model_scanner.py index 3d93145f..eef7cb49 100644 --- a/tests/services/test_model_scanner.py +++ b/tests/services/test_model_scanner.py @@ -1210,3 +1210,111 @@ async def test_bulk_delete_cancelled_after_one_staged_batch_present( assert not first.exists() # The second file was never touched. 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)