From d80c6ff581249367bada241299e4d52094442fd7 Mon Sep 17 00:00:00 2001 From: Robin Huang Date: Fri, 9 Aug 2024 18:07:28 -0700 Subject: [PATCH] Add GET /internal/files. Retrieves list of files in models, output, and user directories. --- api_server/__init__.py | 0 api_server/routes/__init__.py | 0 .../routes}/internal/README.md | 0 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 | 12 +++ api_server/utils/file_operations.py | 42 +++++++++ routes/__init__.py | 1 - routes/internal/internal_routes.py | 65 ------------- server.py | 2 +- tests-unit/conftest.py | 6 ++ .../routes_test/internal_routes_test.py | 94 +++++++++++++++++++ .../services_test/file_service_test.py | 53 +++++++++++ .../utils_test/file_operations_test.py | 42 +++++++++ 15 files changed, 290 insertions(+), 67 deletions(-) create mode 100644 api_server/__init__.py create mode 100644 api_server/routes/__init__.py rename {routes => api_server/routes}/internal/README.md (100%) 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 delete mode 100644 routes/__init__.py delete mode 100644 routes/internal/internal_routes.py create mode 100644 tests-unit/conftest.py create mode 100644 tests-unit/server_test/routes_test/internal_routes_test.py create mode 100644 tests-unit/server_test/services_test/file_service_test.py create mode 100644 tests-unit/server_test/utils_test/file_operations_test.py diff --git a/api_server/__init__.py b/api_server/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/api_server/routes/__init__.py b/api_server/routes/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/routes/internal/README.md b/api_server/routes/internal/README.md similarity index 100% rename from routes/internal/README.md rename to api_server/routes/internal/README.md diff --git a/api_server/routes/internal/__init__.py b/api_server/routes/internal/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/api_server/routes/internal/internal_routes.py b/api_server/routes/internal/internal_routes.py new file mode 100644 index 000000000..2234b3113 --- /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 000000000..e69de29bb diff --git a/api_server/services/file_service.py b/api_server/services/file_service.py new file mode 100644 index 000000000..f56059755 --- /dev/null +++ b/api_server/services/file_service.py @@ -0,0 +1,12 @@ +from api_server.utils.file_operations import FileSystemOperations + +class FileService: + def __init__(self, allowed_directories, file_system_ops=None): + self.allowed_directories = allowed_directories + self.file_system_ops = file_system_ops or FileSystemOperations() + + def list_files(self, directory_key): + if directory_key not in self.allowed_directories: + raise ValueError("Invalid directory key") + directory_path = 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 000000000..b82411ae8 --- /dev/null +++ b/api_server/utils/file_operations.py @@ -0,0 +1,42 @@ +import os +from typing import List, Union, TypedDict, Literal, 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/routes/__init__.py b/routes/__init__.py deleted file mode 100644 index 59ed7f41a..000000000 --- a/routes/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from routes.internal.internal_routes import InternalRoutes \ No newline at end of file diff --git a/routes/internal/internal_routes.py b/routes/internal/internal_routes.py deleted file mode 100644 index cbc61d4d7..000000000 --- a/routes/internal/internal_routes.py +++ /dev/null @@ -1,65 +0,0 @@ -from aiohttp import web -from typing import Optional, Dict, List -from folder_paths import models_dir, user_directory, output_directory -import os - -class InternalRoutes: - def __init__(self): - self.routes: web.RouteTableDef = web.RouteTableDef() - self._app: Optional[web.Application] = None - self.allowed_directories: Dict[str, str] = { - "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', '') - if directory_key not in self.allowed_directories: - return web.json_response({"error": "Invalid directory key"}, status=400) - - directory_path = self.allowed_directories[directory_key] - - try: - file_list = self._get_file_list(directory_path) - return web.json_response({"files": file_list}) - except FileNotFoundError: - return web.json_response({"error": "Directory not found"}, status=404) - except Exception as e: - return web.json_response({"error": str(e)}, status=500) - - def _get_file_list(self, directory: str) -> List[Dict[str, Optional[str]]]: - if not os.path.exists(directory): - raise FileNotFoundError(f"Directory not found: {directory}") - - file_list = [] - 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", - "size": None - }) - return file_list - - - 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/server.py b/server.py index 0c67974e5..ef89cc64f 100644 --- a/server.py +++ b/server.py @@ -29,7 +29,7 @@ from app.frontend_management import FrontendManager from app.user_manager import UserManager from model_filemanager import download_model, DownloadModelStatus from typing import Optional -from routes import InternalRoutes +from api_server.routes.internal.internal_routes import InternalRoutes class BinaryEventTypes: diff --git a/tests-unit/conftest.py b/tests-unit/conftest.py new file mode 100644 index 000000000..e37c19a5d --- /dev/null +++ b/tests-unit/conftest.py @@ -0,0 +1,6 @@ +import sys +import os + +# Add the project root directory to the Python path +project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '..')) +sys.path.insert(0, project_root) \ No newline at end of file diff --git a/tests-unit/server_test/routes_test/internal_routes_test.py b/tests-unit/server_test/routes_test/internal_routes_test.py new file mode 100644 index 000000000..1f17748ad --- /dev/null +++ b/tests-unit/server_test/routes_test/internal_routes_test.py @@ -0,0 +1,94 @@ +import pytest +from aiohttp import web +from unittest.mock import MagicMock, patch +from api_server.routes.internal.internal_routes import InternalRoutes +from folder_paths import models_dir, user_directory, output_directory + + +@pytest.fixture +def internal_routes(): + return InternalRoutes() + +@pytest.fixture +async def client(aiohttp_client, internal_routes): + app = internal_routes.get_app() + return await aiohttp_client(app) + +@pytest.mark.asyncio +async def test_list_files_valid_directory(client, 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) + + 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 + +@pytest.mark.asyncio +async def test_list_files_invalid_directory(client, internal_routes): + internal_routes.file_service.list_files = MagicMock(side_effect=ValueError("Invalid directory key")) + + 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(client, internal_routes): + internal_routes.file_service.list_files = MagicMock(side_effect=Exception("Unexpected error")) + + 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(client, internal_routes): + mock_file_list = [] + internal_routes.file_service.list_files = MagicMock(return_value=mock_file_list) + + resp = await client.get('/files') + assert resp.status == 200 + data = await resp.json() + assert 'files' in data + assert len(data['files']) == 0 + +@patch('server.routes.internal_routes.FileService') +def test_file_service_initialization(mock_file_service): + InternalRoutes() + mock_file_service.assert_called_once_with({ + "models": models_dir, + "user": user_directory, + "output": output_directory + }) + +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 + +# Additional test to check if routes are added to the application +@pytest.mark.asyncio +async def test_routes_added_to_app(aiohttp_client, internal_routes): + app = internal_routes.get_app() + client = await aiohttp_client(app) + + # This will raise an exception if the route doesn't exist + await client.get('/files') \ No newline at end of file diff --git a/tests-unit/server_test/services_test/file_service_test.py b/tests-unit/server_test/services_test/file_service_test.py new file mode 100644 index 000000000..134c668dd --- /dev/null +++ b/tests-unit/server_test/services_test/file_service_test.py @@ -0,0 +1,53 @@ +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): + 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_test/utils_test/file_operations_test.py b/tests-unit/server_test/utils_test/file_operations_test.py new file mode 100644 index 000000000..5a2a83713 --- /dev/null +++ b/tests-unit/server_test/utils_test/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