From f07e5bb5225ff8b9706866a7a5105c987c92b330 Mon Sep 17 00:00:00 2001 From: Robin Huang Date: Tue, 20 Aug 2024 22:25:06 -0700 Subject: [PATCH] Add GET /internal/files. (#4295) * Create internal route table. * List files. * Add GET /internal/files. Retrieves list of files in models, output, and user directories. * Refactor file names. * Use typing_extensions for Python 3.8 * Fix tests. * Remove print statements. * Update README. * Add output and user to valid directory test. * Add missing type hints. --- api_server/__init__.py | 0 api_server/routes/__init__.py | 0 api_server/routes/internal/README.md | 3 + api_server/routes/internal/__init__.py | 0 api_server/routes/internal/internal_routes.py | 40 ++++++ api_server/services/__init__.py | 0 api_server/services/file_service.py | 13 ++ api_server/utils/file_operations.py | 42 +++++++ server.py | 4 + .../server/routes/internal_routes_test.py | 115 ++++++++++++++++++ .../server/services/file_service_test.py | 54 ++++++++ .../server/utils/file_operations_test.py | 42 +++++++ 12 files changed, 313 insertions(+) create mode 100644 api_server/__init__.py create mode 100644 api_server/routes/__init__.py create mode 100644 api_server/routes/internal/README.md create mode 100644 api_server/routes/internal/__init__.py create mode 100644 api_server/routes/internal/internal_routes.py create mode 100644 api_server/services/__init__.py create mode 100644 api_server/services/file_service.py create mode 100644 api_server/utils/file_operations.py create mode 100644 tests-unit/server/routes/internal_routes_test.py create mode 100644 tests-unit/server/services/file_service_test.py create mode 100644 tests-unit/server/utils/file_operations_test.py diff --git a/api_server/__init__.py b/api_server/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/api_server/routes/__init__.py b/api_server/routes/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/api_server/routes/internal/README.md b/api_server/routes/internal/README.md new file mode 100644 index 00000000..35330c36 --- /dev/null +++ b/api_server/routes/internal/README.md @@ -0,0 +1,3 @@ +# ComfyUI Internal Routes + +All routes under the `/internal` path are designated for **internal use by ComfyUI only**. These routes are not intended for use by external applications may change at any time without notice. diff --git a/api_server/routes/internal/__init__.py b/api_server/routes/internal/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/api_server/routes/internal/internal_routes.py b/api_server/routes/internal/internal_routes.py new file mode 100644 index 00000000..2234b311 --- /dev/null +++ b/api_server/routes/internal/internal_routes.py @@ -0,0 +1,40 @@ +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 + +class InternalRoutes: + ''' + The top level web router for internal routes: /internal/* + The endpoints here should NOT be depended upon. It is for ComfyUI frontend use only. + Check README.md for more information. + + ''' + def __init__(self): + self.routes: web.RouteTableDef = web.RouteTableDef() + self._app: Optional[web.Application] = None + self.file_service = FileService({ + "models": models_dir, + "user": user_directory, + "output": output_directory + }) + + def setup_routes(self): + @self.routes.get('/files') + async def list_files(request): + directory_key = request.query.get('directory', '') + try: + file_list = self.file_service.list_files(directory_key) + return web.json_response({"files": file_list}) + except ValueError as e: + return web.json_response({"error": str(e)}, status=400) + except Exception as e: + return web.json_response({"error": str(e)}, status=500) + + + def get_app(self): + if self._app is None: + self._app = web.Application() + self.setup_routes() + self._app.add_routes(self.routes) + return self._app diff --git a/api_server/services/__init__.py b/api_server/services/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/api_server/services/file_service.py b/api_server/services/file_service.py new file mode 100644 index 00000000..39457108 --- /dev/null +++ b/api_server/services/file_service.py @@ -0,0 +1,13 @@ +from typing import Dict, List, Optional +from api_server.utils.file_operations import FileSystemOperations, FileSystemItem + +class FileService: + def __init__(self, allowed_directories: Dict[str, str], file_system_ops: Optional[FileSystemOperations] = None): + self.allowed_directories: Dict[str, str] = allowed_directories + self.file_system_ops: FileSystemOperations = file_system_ops or FileSystemOperations() + + def list_files(self, directory_key: str) -> List[FileSystemItem]: + if directory_key not in self.allowed_directories: + raise ValueError("Invalid directory key") + directory_path: str = self.allowed_directories[directory_key] + return self.file_system_ops.walk_directory(directory_path) \ No newline at end of file diff --git a/api_server/utils/file_operations.py b/api_server/utils/file_operations.py new file mode 100644 index 00000000..ef1bf999 --- /dev/null +++ b/api_server/utils/file_operations.py @@ -0,0 +1,42 @@ +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]: + file_list: List[FileSystemItem] = [] + for root, dirs, files in os.walk(directory): + for name in files: + file_path = os.path.join(root, name) + relative_path = os.path.relpath(file_path, directory) + file_list.append({ + "name": name, + "path": relative_path, + "type": "file", + "size": os.path.getsize(file_path) + }) + for name in dirs: + dir_path = os.path.join(root, name) + relative_path = os.path.relpath(dir_path, directory) + file_list.append({ + "name": name, + "path": relative_path, + "type": "directory" + }) + return file_list \ No newline at end of file diff --git a/server.py b/server.py index 23ebdb7b..edb4a8ca 100644 --- a/server.py +++ b/server.py @@ -29,6 +29,8 @@ from app.frontend_management import FrontendManager from app.user_manager import UserManager from model_filemanager import download_model, DownloadModelStatus from typing import Optional +from api_server.routes.internal.internal_routes import InternalRoutes + class BinaryEventTypes: PREVIEW_IMAGE = 1 @@ -72,6 +74,7 @@ class PromptServer(): mimetypes.types_map['.js'] = 'application/javascript; charset=utf-8' self.user_manager = UserManager() + self.internal_routes = InternalRoutes() self.supports = ["custom_nodes_from_web"] self.prompt_queue = None self.loop = loop @@ -602,6 +605,7 @@ class PromptServer(): def add_routes(self): self.user_manager.add_routes(self.routes) + self.app.add_subapp('/internal', self.internal_routes.get_app()) # Prefix every route with /api for easier matching for delegation. # This is very useful for frontend dev server, which need to forward diff --git a/tests-unit/server/routes/internal_routes_test.py b/tests-unit/server/routes/internal_routes_test.py new file mode 100644 index 00000000..2d2b43bd --- /dev/null +++ b/tests-unit/server/routes/internal_routes_test.py @@ -0,0 +1,115 @@ +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 + + +@pytest.fixture +def internal_routes(): + return InternalRoutes() + +@pytest.fixture +def aiohttp_client_factory(aiohttp_client, internal_routes): + async def _get_client(): + app = internal_routes.get_app() + return await aiohttp_client(app) + return _get_client + +@pytest.mark.asyncio +async def test_list_files_valid_directory(aiohttp_client_factory, internal_routes): + mock_file_list = [ + {"name": "file1.txt", "path": "file1.txt", "type": "file", "size": 100}, + {"name": "dir1", "path": "dir1", "type": "directory"} + ] + internal_routes.file_service.list_files = MagicMock(return_value=mock_file_list) + client = await aiohttp_client_factory() + resp = await client.get('/files?directory=models') + assert resp.status == 200 + data = await resp.json() + assert 'files' in data + assert len(data['files']) == 2 + assert data['files'] == mock_file_list + + # Check other valid directories + resp = await client.get('/files?directory=user') + assert resp.status == 200 + resp = await client.get('/files?directory=output') + assert resp.status == 200 + +@pytest.mark.asyncio +async def test_list_files_invalid_directory(aiohttp_client_factory, internal_routes): + internal_routes.file_service.list_files = MagicMock(side_effect=ValueError("Invalid directory key")) + client = await aiohttp_client_factory() + resp = await client.get('/files?directory=invalid') + assert resp.status == 400 + data = await resp.json() + assert 'error' in data + assert data['error'] == "Invalid directory key" + +@pytest.mark.asyncio +async def test_list_files_exception(aiohttp_client_factory, internal_routes): + internal_routes.file_service.list_files = MagicMock(side_effect=Exception("Unexpected error")) + client = await aiohttp_client_factory() + resp = await client.get('/files?directory=models') + assert resp.status == 500 + data = await resp.json() + assert 'error' in data + assert data['error'] == "Unexpected error" + +@pytest.mark.asyncio +async def test_list_files_no_directory_param(aiohttp_client_factory, internal_routes): + mock_file_list = [] + internal_routes.file_service.list_files = MagicMock(return_value=mock_file_list) + client = await aiohttp_client_factory() + resp = await client.get('/files') + assert resp.status == 200 + data = await resp.json() + assert 'files' in data + assert len(data['files']) == 0 + +def test_setup_routes(internal_routes): + internal_routes.setup_routes() + routes = internal_routes.routes + assert any(route.method == 'GET' and str(route.path) == '/files' for route in routes) + +def test_get_app(internal_routes): + app = internal_routes.get_app() + assert isinstance(app, web.Application) + assert internal_routes._app is not None + +def test_get_app_reuse(internal_routes): + app1 = internal_routes.get_app() + app2 = internal_routes.get_app() + assert app1 is app2 + +@pytest.mark.asyncio +async def test_routes_added_to_app(aiohttp_client_factory, internal_routes): + client = await aiohttp_client_factory() + try: + resp = await client.get('/files') + print(f"Response received: status {resp.status}") + except Exception as e: + print(f"Exception occurred during GET request: {e}") + raise + + assert resp.status != 404, "Route /files does not exist" + +@pytest.mark.asyncio +async def test_file_service_initialization(): + with patch('api_server.routes.internal.internal_routes.FileService') as MockFileService: + # Create a mock instance + mock_file_service_instance = MagicMock(spec=FileService) + MockFileService.return_value = mock_file_service_instance + internal_routes = InternalRoutes() + + # Check if FileService was initialized with the correct parameters + MockFileService.assert_called_once_with({ + "models": models_dir, + "user": user_directory, + "output": output_directory + }) + + # Verify that the file_service attribute of InternalRoutes is set + assert internal_routes.file_service == mock_file_service_instance \ No newline at end of file diff --git a/tests-unit/server/services/file_service_test.py b/tests-unit/server/services/file_service_test.py new file mode 100644 index 00000000..5650452a --- /dev/null +++ b/tests-unit/server/services/file_service_test.py @@ -0,0 +1,54 @@ +import pytest +from unittest.mock import MagicMock +from 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 = { + "models": "/path/to/models", + "user": "/path/to/user", + "output": "/path/to/output" + } + 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 diff --git a/tests-unit/server/utils/file_operations_test.py b/tests-unit/server/utils/file_operations_test.py new file mode 100644 index 00000000..5a2a8371 --- /dev/null +++ b/tests-unit/server/utils/file_operations_test.py @@ -0,0 +1,42 @@ +import pytest +from typing import List +from api_server.utils.file_operations import FileSystemOperations, FileSystemItem, is_file_info + +@pytest.fixture +def temp_directory(tmp_path): + # Create a temporary directory structure + dir1 = tmp_path / "dir1" + dir2 = tmp_path / "dir2" + dir1.mkdir() + dir2.mkdir() + (dir1 / "file1.txt").write_text("content1") + (dir2 / "file2.txt").write_text("content2") + (tmp_path / "file3.txt").write_text("content3") + return tmp_path + +def test_walk_directory(temp_directory): + result: List[FileSystemItem] = FileSystemOperations.walk_directory(str(temp_directory)) + + assert len(result) == 5 # 2 directories and 3 files + + files = [item for item in result if item['type'] == 'file'] + dirs = [item for item in result if item['type'] == 'directory'] + + assert len(files) == 3 + assert len(dirs) == 2 + + file_names = {file['name'] for file in files} + assert file_names == {'file1.txt', 'file2.txt', 'file3.txt'} + + dir_names = {dir['name'] for dir in dirs} + assert dir_names == {'dir1', 'dir2'} + +def test_walk_directory_empty(tmp_path): + result = FileSystemOperations.walk_directory(str(tmp_path)) + assert len(result) == 0 + +def test_walk_directory_file_size(temp_directory): + result: List[FileSystemItem] = FileSystemOperations.walk_directory(str(temp_directory)) + files = [item for item in result if is_file_info(item)] + for file in files: + assert file['size'] > 0 # Assuming all files have some content