refactor into custom_node_manager, add test

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)])

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]] = {}
 

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)])

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"