refactor: move cache-controls to /middleware

We can test the code without excessively mocking server dependencies

middleware/__init__.py

@@ -0,0 +1 @@
+"""Server middleware modules"""

middleware/cache_middleware.py

@@ -0,0 +1,46 @@
+"""Cache control middleware for ComfyUI server"""
+
+from aiohttp import web
+from typing import Callable, Awaitable
+
+# Time in seconds
+ONE_HOUR: int = 3600
+ONE_DAY: int = 86400
+IMG_EXTENSIONS = (
+    ".jpg",
+    ".jpeg",
+    ".png",
+    ".ppm",
+    ".bmp",
+    ".pgm",
+    ".tif",
+    ".tiff",
+    ".webp",
+)
+
+
+@web.middleware
+async def cache_control(
+    request: web.Request, handler: Callable[[web.Request], Awaitable[web.Response]]
+) -> web.Response:
+    """Cache control middleware that sets appropriate cache headers based on file type and response status"""
+    response: web.Response = await handler(request)
+
+    if (
+        request.path.endswith(".js")
+        or request.path.endswith(".css")
+        or request.path.endswith("index.json")
+    ):
+        response.headers.setdefault("Cache-Control", "no-cache")
+    elif request.path.lower().endswith(IMG_EXTENSIONS):
+        if response.status == 404:
+            response.headers.setdefault("Cache-Control", f"public, max-age={ONE_HOUR}")
+        elif response.status in (200, 201, 202, 203, 204, 205, 206, 301, 308):
+            # Success responses and permanent redirects - cache for 1 day
+            response.headers.setdefault("Cache-Control", f"public, max-age={ONE_DAY}")
+        elif response.status in (302, 303, 307):
+            # Temporary redirects - no cache
+            response.headers.setdefault("Cache-Control", "no-cache")
+        # Note: 304 Not Modified falls through - no cache headers set
+
+    return response

server.py

@@ -39,10 +39,8 @@ from typing import Optional, Union, Callable, Awaitable
 from api_server.routes.internal.internal_routes import InternalRoutes
 from protocol import BinaryEventTypes
 
-# Time in seconds
-ONE_HOUR: int = 3600
-ONE_DAY: int = 86400
-IMG_EXTENSIONS = ('.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp')
+# Import cache control middleware
+from middleware.cache_middleware import cache_control, ONE_HOUR, ONE_DAY, IMG_EXTENSIONS
 
 async def send_socket_catch_exception(function, message):
     try:
@@ -51,23 +49,6 @@ async def send_socket_catch_exception(function, message):
         logging.warning("send error: {}".format(err))
 
 @web.middleware
-async def cache_control(request: web.Request, handler: Callable[[web.Request], Awaitable[web.Response]]) -> web.Response:
-    response: web.Response = await handler(request)
-
-    if request.path.endswith('.js') or request.path.endswith('.css') or request.path.endswith('index.json'):
-        response.headers.setdefault('Cache-Control', 'no-cache')
-    elif request.path.lower().endswith(IMG_EXTENSIONS):
-        if response.status == 404:
-            response.headers.setdefault('Cache-Control', f"public, max-age={ONE_HOUR}")
-        elif response.status in (200, 201, 202, 203, 204, 205, 206, 301, 308):
-            # Success responses and permanent redirects - cache for 1 day
-            response.headers.setdefault('Cache-Control', f"public, max-age={ONE_DAY}")
-        elif response.status in (302, 303, 307):
-            # Temporary redirects - no cache
-            response.headers.setdefault('Cache-Control', 'no-cache')
-        # Note: 304 Not Modified falls through - no cache headers set
-
-    return response
 
 @web.middleware
 async def compress_body(request: web.Request, handler):

tests-unit/server_test/test_cache_control.py

@@ -3,20 +3,11 @@
 import pytest
 from aiohttp import web
 from aiohttp.test_utils import make_mocked_request
-from unittest.mock import patch
+
+from middleware.cache_middleware import cache_control, ONE_HOUR, ONE_DAY, IMG_EXTENSIONS
 
 pytestmark = pytest.mark.asyncio  # Apply asyncio mark to all tests
 
-# Mock the problematic imports before importing server
-with (
-    patch("app.frontend_management.FrontendManager"),
-    patch("utils.install_util.get_missing_requirements_message"),
-    patch("utils.install_util.requirements_path"),
-    patch("comfy.model_management.get_torch_device"),
-    patch("comfy.model_management.get_total_memory", return_value=1024 * 1024 * 1024),
-):
-    from server import cache_control, ONE_HOUR, ONE_DAY, IMG_EXTENSIONS
-
 
 class TestCacheControl:
     """Test cache control middleware functionality"""