From fc14655ac4fa1c303f9d6d41c0e740bc88c9781c Mon Sep 17 00:00:00 2001 From: bezo97 Date: Wed, 25 Dec 2024 15:45:08 +0100 Subject: [PATCH] refactor into custom_node_manager, add test --- app/custom_node_manager.py | 30 ++++++++++++++ folder_paths.py | 1 - server.py | 19 ++------- .../app_test/custom_node_manager_test.py | 41 +++++++++++++++++++ 4 files changed, 74 insertions(+), 17 deletions(-) create mode 100644 app/custom_node_manager.py create mode 100644 tests-unit/app_test/custom_node_manager_test.py diff --git a/app/custom_node_manager.py b/app/custom_node_manager.py new file mode 100644 index 000000000..f315d3ca9 --- /dev/null +++ b/app/custom_node_manager.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +import os +import folder_paths +import glob +from aiohttp import web + +class CustomNodeManager: + """ + Placeholder to refactor the custom node management features from ComfyUI-Manager. + Currently it only contains the custom workflow templates feature. + """ + def add_routes(self, routes, webapp, loadedModules): + + @routes.get("/workflow_templates") + async def get_workflow_templates(request): + """Returns a web response that contains the map of custom_nodes names and their associated workflow templates. The ones without templates are omitted.""" + files = glob.glob(os.path.join(folder_paths.get_folder_paths("custom_nodes")[0], '*/example_workflows/*.json')) + workflow_templates_dict = {} # custom_nodes folder name -> example workflow names + for file in files: + custom_nodes_name = os.path.basename(os.path.dirname(os.path.dirname(file))) + workflow_name = os.path.splitext(os.path.basename(file))[0] + workflow_templates_dict.setdefault(custom_nodes_name, []).append(workflow_name) + return web.json_response(workflow_templates_dict) + + # Serve workflow templates from custom nodes. + for module_name, module_dir in loadedModules: + workflows_dir = os.path.join(module_dir, 'example_workflows') + if os.path.exists(workflows_dir): + webapp.add_routes([web.static('/api/workflow_templates/' + module_name, workflows_dir)]) \ No newline at end of file diff --git a/folder_paths.py b/folder_paths.py index 8494080a9..796affa5f 100644 --- a/folder_paths.py +++ b/folder_paths.py @@ -43,7 +43,6 @@ output_directory = os.path.join(base_path, "output") temp_directory = os.path.join(base_path, "temp") input_directory = os.path.join(base_path, "input") user_directory = os.path.join(base_path, "user") -custom_nodes_directory = os.path.join(base_path, "custom_nodes") filename_list_cache: dict[str, tuple[list[str], dict[str, float], float]] = {} diff --git a/server.py b/server.py index c868dffab..85fb5732a 100644 --- a/server.py +++ b/server.py @@ -30,6 +30,7 @@ import node_helpers from app.frontend_management import FrontendManager from app.user_manager import UserManager from app.model_manager import ModelFileManager +from app.custom_node_manager import CustomNodeManager from typing import Optional from api_server.routes.internal.internal_routes import InternalRoutes @@ -153,6 +154,7 @@ class PromptServer(): self.user_manager = UserManager() self.model_file_manager = ModelFileManager() + self.custom_node_manager = CustomNodeManager() self.internal_routes = InternalRoutes(self) self.supports = ["custom_nodes_from_web"] self.prompt_queue = None @@ -251,16 +253,6 @@ class PromptServer(): return web.json_response(extensions) - @routes.get("/workflow_templates") - async def get_workflow_templates(request): - files = glob.glob(os.path.join(folder_paths.custom_nodes_directory, '*/example_workflows/*.json')) - workflow_templates_dict = {} # custom_nodes folder name -> example workflow names - for file in files: - custom_nodes_name = os.path.basename(os.path.dirname(os.path.dirname(file))) - workflow_name = os.path.splitext(os.path.basename(file))[0] - workflow_templates_dict.setdefault(custom_nodes_name, []).append(workflow_name) - return web.json_response(workflow_templates_dict) - def get_dir_by_type(dir_type): if dir_type is None: dir_type = "input" @@ -707,6 +699,7 @@ class PromptServer(): def add_routes(self): self.user_manager.add_routes(self.routes) self.model_file_manager.add_routes(self.routes) + self.custom_node_manager.add_routes(self.routes, self.app, nodes.LOADED_MODULE_DIRS.items()) self.app.add_subapp('/internal', self.internal_routes.get_app()) # Prefix every route with /api for easier matching for delegation. @@ -723,12 +716,6 @@ class PromptServer(): self.app.add_routes(api_routes) self.app.add_routes(self.routes) - # Add routes for workflow templates in custom nodes. - for module_name, module_dir in nodes.LOADED_MODULE_DIRS.items(): - workflows_dir = os.path.join(module_dir, 'example_workflows') - if os.path.exists(workflows_dir): - self.app.add_routes([web.static('/workflow_templates/' + module_name, workflows_dir)]) - # Add routes from web extensions. for name, dir in nodes.EXTENSION_WEB_DIRS.items(): self.app.add_routes([web.static('/extensions/' + name, dir)]) diff --git a/tests-unit/app_test/custom_node_manager_test.py b/tests-unit/app_test/custom_node_manager_test.py new file mode 100644 index 000000000..8df9d5225 --- /dev/null +++ b/tests-unit/app_test/custom_node_manager_test.py @@ -0,0 +1,41 @@ +import pytest +import json +from aiohttp import web +from unittest.mock import patch +from app.custom_node_manager import CustomNodeManager + +pytestmark = ( + pytest.mark.asyncio +) # This applies the asyncio mark to all test functions in the module + +@pytest.fixture +def custom_node_manager(): + return CustomNodeManager() + +@pytest.fixture +def app(custom_node_manager): + app = web.Application() + routes = web.RouteTableDef() + custom_node_manager.add_routes(routes, app, [("ComfyUI-TestExtension1", "ComfyUI-TestExtension1")]) + app.add_routes(routes) + return app + +async def test_get_workflow_templates(aiohttp_client, app, tmp_path): + client = await aiohttp_client(app) + # Setup temporary custom nodes file structure with 1 workflow file + custom_nodes_dir = tmp_path / "custom_nodes" + example_workflows_dir = custom_nodes_dir / "ComfyUI-TestExtension1" / "example_workflows" + example_workflows_dir.mkdir(parents=True) + template_file = example_workflows_dir / "workflow1.json" + template_file.write_text('') + + with patch('folder_paths.folder_names_and_paths', { + 'custom_nodes': ([str(custom_nodes_dir)], None) + }): + response = await client.get('/workflow_templates') + assert response.status == 200 + workflows_dict = await response.json() + assert isinstance(workflows_dict, dict) + assert "ComfyUI-TestExtension1" in workflows_dict + assert isinstance(workflows_dict["ComfyUI-TestExtension1"], list) + assert workflows_dict["ComfyUI-TestExtension1"][0] == "workflow1"