diff --git a/comfy/api_server/routes/internal/internal_routes.py b/comfy/api_server/routes/internal/internal_routes.py index 2234b3113..0b1216ebd 100644 --- a/comfy/api_server/routes/internal/internal_routes.py +++ b/comfy/api_server/routes/internal/internal_routes.py @@ -1,7 +1,7 @@ from aiohttp import web from typing import Optional -from folder_paths import models_dir, user_directory, output_directory -from api_server.services.file_service import FileService +from ....cmd.folder_paths import models_dir, user_directory, output_directory +from ...services.file_service import FileService class InternalRoutes: ''' diff --git a/comfy/api_server/services/file_service.py b/comfy/api_server/services/file_service.py index 394571084..09eea6c09 100644 --- a/comfy/api_server/services/file_service.py +++ b/comfy/api_server/services/file_service.py @@ -1,5 +1,5 @@ from typing import Dict, List, Optional -from api_server.utils.file_operations import FileSystemOperations, FileSystemItem +from ..utils.file_operations import FileSystemOperations, FileSystemItem class FileService: def __init__(self, allowed_directories: Dict[str, str], file_system_ops: Optional[FileSystemOperations] = None): diff --git a/comfy/api_server/utils/__init__.py b/comfy/api_server/utils/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/comfy/api_server/utils/file_operations.py b/comfy/api_server/utils/file_operations.py index ef1bf999e..2c76a6cdd 100644 --- a/comfy/api_server/utils/file_operations.py +++ b/comfy/api_server/utils/file_operations.py @@ -1,22 +1,29 @@ import os from typing import List, Union, TypedDict, Literal + from typing_extensions import TypeGuard + + class FileInfo(TypedDict): name: str path: str type: Literal["file"] size: int + class DirectoryInfo(TypedDict): name: str path: str type: Literal["directory"] + FileSystemItem = Union[FileInfo, DirectoryInfo] + def is_file_info(item: FileSystemItem) -> TypeGuard[FileInfo]: return item["type"] == "file" + class FileSystemOperations: @staticmethod def walk_directory(directory: str) -> List[FileSystemItem]: @@ -39,4 +46,4 @@ class FileSystemOperations: "path": relative_path, "type": "directory" }) - return file_list \ No newline at end of file + return file_list diff --git a/comfy/cmd/server.py b/comfy/cmd/server.py index d5a63ddf5..a9ddb9cd1 100644 --- a/comfy/cmd/server.py +++ b/comfy/cmd/server.py @@ -25,7 +25,7 @@ from aiohttp import web from can_ada import URL, parse as urlparse # pylint: disable=no-name-in-module from typing_extensions import NamedTuple -from api_server.routes.internal.internal_routes import InternalRoutes +from ..api_server.routes.internal.internal_routes import InternalRoutes from ..model_filemanager import download_model, DownloadModelStatus from .latent_preview_image_encoding import encode_preview_image from .. import interruption diff --git a/tests/unit/server/__init__.py b/tests/unit/server/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit/server/routes/__init__.py b/tests/unit/server/routes/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit/server/routes/internal_routes_test.py b/tests/unit/server/routes/internal_routes_test.py index 2d2b43bd6..f3b1918ea 100644 --- a/tests/unit/server/routes/internal_routes_test.py +++ b/tests/unit/server/routes/internal_routes_test.py @@ -1,9 +1,9 @@ import pytest from aiohttp import web from unittest.mock import MagicMock, patch -from api_server.routes.internal.internal_routes import InternalRoutes -from api_server.services.file_service import FileService -from folder_paths import models_dir, user_directory, output_directory +from comfy.api_server.routes.internal.internal_routes import InternalRoutes +from comfy.api_server.services.file_service import FileService +from comfy.cmd.folder_paths import models_dir, user_directory, output_directory @pytest.fixture diff --git a/tests/unit/server/services/__init__.py b/tests/unit/server/services/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit/server/services/file_service_test.py b/tests/unit/server/services/file_service_test.py index 5650452a3..c8aca52b5 100644 --- a/tests/unit/server/services/file_service_test.py +++ b/tests/unit/server/services/file_service_test.py @@ -1,11 +1,15 @@ -import pytest from unittest.mock import MagicMock -from api_server.services.file_service import FileService + +import pytest + +from comfy.api_server.services.file_service import FileService + @pytest.fixture def mock_file_system_ops(): return MagicMock() + @pytest.fixture def file_service(mock_file_system_ops): allowed_directories = { @@ -15,40 +19,44 @@ def file_service(mock_file_system_ops): } return FileService(allowed_directories, file_system_ops=mock_file_system_ops) + def test_list_files_valid_directory(file_service, mock_file_system_ops): mock_file_system_ops.walk_directory.return_value = [ {"name": "file1.txt", "path": "file1.txt", "type": "file", "size": 100}, {"name": "dir1", "path": "dir1", "type": "directory"} ] - + result = file_service.list_files("models") - + assert len(result) == 2 assert result[0]["name"] == "file1.txt" assert result[1]["name"] == "dir1" mock_file_system_ops.walk_directory.assert_called_once_with("/path/to/models") + def test_list_files_invalid_directory(file_service): # Does not support walking directories outside of the allowed directories with pytest.raises(ValueError, match="Invalid directory key"): file_service.list_files("invalid_key") + def test_list_files_empty_directory(file_service, mock_file_system_ops): mock_file_system_ops.walk_directory.return_value = [] - + result = file_service.list_files("models") - + assert len(result) == 0 mock_file_system_ops.walk_directory.assert_called_once_with("/path/to/models") + @pytest.mark.parametrize("directory_key", ["models", "user", "output"]) def test_list_files_all_allowed_directories(file_service, mock_file_system_ops, directory_key): mock_file_system_ops.walk_directory.return_value = [ {"name": f"file_{directory_key}.txt", "path": f"file_{directory_key}.txt", "type": "file", "size": 100} ] - + result = file_service.list_files(directory_key) - + assert len(result) == 1 assert result[0]["name"] == f"file_{directory_key}.txt" - mock_file_system_ops.walk_directory.assert_called_once_with(f"/path/to/{directory_key}") \ No newline at end of file + mock_file_system_ops.walk_directory.assert_called_once_with(f"/path/to/{directory_key}") diff --git a/tests/unit/server/utils/__init__.py b/tests/unit/server/utils/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/unit/server/utils/file_operations_test.py b/tests/unit/server/utils/file_operations_test.py index 5a2a83713..14cac65d1 100644 --- a/tests/unit/server/utils/file_operations_test.py +++ b/tests/unit/server/utils/file_operations_test.py @@ -1,6 +1,6 @@ import pytest from typing import List -from api_server.utils.file_operations import FileSystemOperations, FileSystemItem, is_file_info +from comfy.api_server.utils.file_operations import FileSystemOperations, FileSystemItem, is_file_info @pytest.fixture def temp_directory(tmp_path):