Merge pull request #461 from willmiao/codex/refactor-example-images-routes

Add regression tests for example images routes

py/routes/example_images_route_registrar.py

@@ -0,0 +1,61 @@
+"""Route registrar for example image endpoints."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable, Iterable, Mapping
+
+from aiohttp import web
+
+
+@dataclass(frozen=True)
+class RouteDefinition:
+    """Declarative configuration for a HTTP route."""
+
+    method: str
+    path: str
+    handler_name: str
+
+
+ROUTE_DEFINITIONS: tuple[RouteDefinition, ...] = (
+    RouteDefinition("POST", "/api/lm/download-example-images", "download_example_images"),
+    RouteDefinition("POST", "/api/lm/import-example-images", "import_example_images"),
+    RouteDefinition("GET", "/api/lm/example-images-status", "get_example_images_status"),
+    RouteDefinition("POST", "/api/lm/pause-example-images", "pause_example_images"),
+    RouteDefinition("POST", "/api/lm/resume-example-images", "resume_example_images"),
+    RouteDefinition("POST", "/api/lm/open-example-images-folder", "open_example_images_folder"),
+    RouteDefinition("GET", "/api/lm/example-image-files", "get_example_image_files"),
+    RouteDefinition("GET", "/api/lm/has-example-images", "has_example_images"),
+    RouteDefinition("POST", "/api/lm/delete-example-image", "delete_example_image"),
+    RouteDefinition("POST", "/api/lm/force-download-example-images", "force_download_example_images"),
+)
+
+
+class ExampleImagesRouteRegistrar:
+    """Bind declarative example image routes to an aiohttp router."""
+
+    _METHOD_MAP = {
+        "GET": "add_get",
+        "POST": "add_post",
+        "PUT": "add_put",
+        "DELETE": "add_delete",
+    }
+
+    def __init__(self, app: web.Application) -> None:
+        self._app = app
+
+    def register_routes(
+        self,
+        handler_lookup: Mapping[str, Callable[[web.Request], object]],
+        *,
+        definitions: Iterable[RouteDefinition] = ROUTE_DEFINITIONS,
+    ) -> None:
+        """Register each route definition using the supplied handlers."""
+
+        for definition in definitions:
+            handler = handler_lookup[definition.handler_name]
+            self._bind_route(definition.method, definition.path, handler)
+
+    def _bind_route(self, method: str, path: str, handler: Callable[[web.Request], object]) -> None:
+        add_method_name = self._METHOD_MAP[method.upper()]
+        add_method = getattr(self._app.router, add_method_name)
+        add_method(path, handler)

py/routes/example_images_routes.py

@@ -1,74 +1,69 @@
+from __future__ import annotations
+
 import logging
+from typing import Callable, Mapping
+
+from aiohttp import web
+
+from .example_images_route_registrar import ExampleImagesRouteRegistrar
+from .handlers.example_images_handlers import (
+    ExampleImagesDownloadHandler,
+    ExampleImagesFileHandler,
+    ExampleImagesHandlerSet,
+    ExampleImagesManagementHandler,
+)
 from ..utils.example_images_download_manager import DownloadManager
-from ..utils.example_images_processor import ExampleImagesProcessor
 from ..utils.example_images_file_manager import ExampleImagesFileManager
-from ..services.websocket_manager import ws_manager
+from ..utils.example_images_processor import ExampleImagesProcessor
 
 logger = logging.getLogger(__name__)
 
+
 class ExampleImagesRoutes:
-    """Routes for example images related functionality"""
+    """Route controller for example image endpoints."""
-    
-    @staticmethod
-    def setup_routes(app):
-        """Register example images routes"""
-        app.router.add_post('/api/lm/download-example-images', ExampleImagesRoutes.download_example_images)
-        app.router.add_post('/api/lm/import-example-images', ExampleImagesRoutes.import_example_images)
-        app.router.add_get('/api/lm/example-images-status', ExampleImagesRoutes.get_example_images_status)
-        app.router.add_post('/api/lm/pause-example-images', ExampleImagesRoutes.pause_example_images)
-        app.router.add_post('/api/lm/resume-example-images', ExampleImagesRoutes.resume_example_images)
-        app.router.add_post('/api/lm/open-example-images-folder', ExampleImagesRoutes.open_example_images_folder)
-        app.router.add_get('/api/lm/example-image-files', ExampleImagesRoutes.get_example_image_files)
-        app.router.add_get('/api/lm/has-example-images', ExampleImagesRoutes.has_example_images)
-        app.router.add_post('/api/lm/delete-example-image', ExampleImagesRoutes.delete_example_image)
-        app.router.add_post('/api/lm/force-download-example-images', ExampleImagesRoutes.force_download_example_images)
 
-    @staticmethod
+    def __init__(
-    async def download_example_images(request):
+        self,
-        """Download example images for models from Civitai"""
+        *,
-        return await DownloadManager.start_download(request)
+        download_manager=DownloadManager,
+        processor=ExampleImagesProcessor,
+        file_manager=ExampleImagesFileManager,
+    ) -> None:
+        self._download_manager = download_manager
+        self._processor = processor
+        self._file_manager = file_manager
+        self._handler_set: ExampleImagesHandlerSet | None = None
+        self._handler_mapping: Mapping[str, Callable[[web.Request], web.StreamResponse]] | None = None
 
-    @staticmethod
+    @classmethod
-    async def get_example_images_status(request):
+    def setup_routes(cls, app: web.Application) -> None:
-        """Get the current status of example images download"""
+        """Register routes on the given aiohttp application using default wiring."""
-        return await DownloadManager.get_status(request)
 
-    @staticmethod
+        controller = cls()
-    async def pause_example_images(request):
+        controller.register(app)
-        """Pause the example images download"""
-        return await DownloadManager.pause_download(request)
 
-    @staticmethod
+    def register(self, app: web.Application) -> None:
-    async def resume_example_images(request):
+        """Bind the controller's handlers to the aiohttp router."""
-        """Resume the example images download"""
-        return await DownloadManager.resume_download(request)
-        
-    @staticmethod
-    async def open_example_images_folder(request):
-        """Open the example images folder for a specific model"""
-        return await ExampleImagesFileManager.open_folder(request)
 
-    @staticmethod
+        registrar = ExampleImagesRouteRegistrar(app)
-    async def get_example_image_files(request):
+        registrar.register_routes(self.to_route_mapping())
-        """Get list of example image files for a specific model"""
-        return await ExampleImagesFileManager.get_files(request)
 
-    @staticmethod
+    def to_route_mapping(self) -> Mapping[str, Callable[[web.Request], web.StreamResponse]]:
-    async def import_example_images(request):
+        """Return the registrar-compatible mapping of handler names to callables."""
-        """Import local example images for a model"""
-        return await ExampleImagesProcessor.import_images(request)
-        
-    @staticmethod
-    async def has_example_images(request):
-        """Check if example images folder exists and is not empty for a model"""
-        return await ExampleImagesFileManager.has_images(request)
 
-    @staticmethod
+        if self._handler_mapping is None:
-    async def delete_example_image(request):
+            handler_set = self._build_handler_set()
-        """Delete a custom example image for a model"""
+            self._handler_set = handler_set
-        return await ExampleImagesProcessor.delete_custom_image(request)
+            self._handler_mapping = handler_set.to_route_mapping()
+        return self._handler_mapping
 
-    @staticmethod
+    def _build_handler_set(self) -> ExampleImagesHandlerSet:
-    async def force_download_example_images(request):
+        logger.debug("Building ExampleImagesHandlerSet with %s, %s, %s", self._download_manager, self._processor, self._file_manager)
-        """Force download example images for specific models"""
+        download_handler = ExampleImagesDownloadHandler(self._download_manager)
-        return await DownloadManager.start_force_download(request)
+        management_handler = ExampleImagesManagementHandler(self._processor)
+        file_handler = ExampleImagesFileHandler(self._file_manager)
+        return ExampleImagesHandlerSet(
+            download=download_handler,
+            management=management_handler,
+            files=file_handler,
+        )

py/routes/handlers/example_images_handlers.py

@@ -0,0 +1,83 @@
+"""Handler set for example image routes."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable, Mapping
+
+from aiohttp import web
+
+
+class ExampleImagesDownloadHandler:
+    """HTTP adapters for download-related example image endpoints."""
+
+    def __init__(self, download_manager) -> None:
+        self._download_manager = download_manager
+
+    async def download_example_images(self, request: web.Request) -> web.StreamResponse:
+        return await self._download_manager.start_download(request)
+
+    async def get_example_images_status(self, request: web.Request) -> web.StreamResponse:
+        return await self._download_manager.get_status(request)
+
+    async def pause_example_images(self, request: web.Request) -> web.StreamResponse:
+        return await self._download_manager.pause_download(request)
+
+    async def resume_example_images(self, request: web.Request) -> web.StreamResponse:
+        return await self._download_manager.resume_download(request)
+
+    async def force_download_example_images(self, request: web.Request) -> web.StreamResponse:
+        return await self._download_manager.start_force_download(request)
+
+
+class ExampleImagesManagementHandler:
+    """HTTP adapters for import/delete endpoints."""
+
+    def __init__(self, processor) -> None:
+        self._processor = processor
+
+    async def import_example_images(self, request: web.Request) -> web.StreamResponse:
+        return await self._processor.import_images(request)
+
+    async def delete_example_image(self, request: web.Request) -> web.StreamResponse:
+        return await self._processor.delete_custom_image(request)
+
+
+class ExampleImagesFileHandler:
+    """HTTP adapters for filesystem-centric endpoints."""
+
+    def __init__(self, file_manager) -> None:
+        self._file_manager = file_manager
+
+    async def open_example_images_folder(self, request: web.Request) -> web.StreamResponse:
+        return await self._file_manager.open_folder(request)
+
+    async def get_example_image_files(self, request: web.Request) -> web.StreamResponse:
+        return await self._file_manager.get_files(request)
+
+    async def has_example_images(self, request: web.Request) -> web.StreamResponse:
+        return await self._file_manager.has_images(request)
+
+
+@dataclass(frozen=True)
+class ExampleImagesHandlerSet:
+    """Aggregate of handlers exposed to the registrar."""
+
+    download: ExampleImagesDownloadHandler
+    management: ExampleImagesManagementHandler
+    files: ExampleImagesFileHandler
+
+    def to_route_mapping(self) -> Mapping[str, Callable[[web.Request], web.StreamResponse]]:
+        """Flatten handler methods into the registrar mapping."""
+
+        return {
+            "download_example_images": self.download.download_example_images,
+            "get_example_images_status": self.download.get_example_images_status,
+            "pause_example_images": self.download.pause_example_images,
+            "resume_example_images": self.download.resume_example_images,
+            "force_download_example_images": self.download.force_download_example_images,
+            "import_example_images": self.management.import_example_images,
+            "delete_example_image": self.management.delete_example_image,
+            "open_example_images_folder": self.files.open_example_images_folder,
+            "get_example_image_files": self.files.get_example_image_files,
+            "has_example_images": self.files.has_example_images,
+        }

tests/routes/test_example_images_routes.py

@@ -0,0 +1,379 @@
+from __future__ import annotations
+
+from contextlib import asynccontextmanager
+from dataclasses import dataclass
+from typing import Any, List, Tuple
+
+from aiohttp import web
+from aiohttp.test_utils import TestClient, TestServer
+import pytest
+
+from py.routes.example_images_route_registrar import ROUTE_DEFINITIONS
+from py.routes.example_images_routes import ExampleImagesRoutes
+from py.routes.handlers.example_images_handlers import (
+    ExampleImagesDownloadHandler,
+    ExampleImagesFileHandler,
+    ExampleImagesHandlerSet,
+    ExampleImagesManagementHandler,
+)
+
+
+@dataclass
+class ExampleImagesHarness:
+    """Container exposing the aiohttp client and stubbed collaborators."""
+
+    client: TestClient
+    download_manager: "StubDownloadManager"
+    processor: "StubExampleImagesProcessor"
+    file_manager: "StubExampleImagesFileManager"
+    controller: ExampleImagesRoutes
+
+
+class StubDownloadManager:
+    def __init__(self) -> None:
+        self.calls: List[Tuple[str, Any]] = []
+
+    async def start_download(self, request: web.Request) -> web.StreamResponse:
+        payload = await request.json()
+        self.calls.append(("start_download", payload))
+        return web.json_response({"operation": "start_download", "payload": payload})
+
+    async def get_status(self, request: web.Request) -> web.StreamResponse:
+        self.calls.append(("get_status", dict(request.query)))
+        return web.json_response({"operation": "get_status"})
+
+    async def pause_download(self, request: web.Request) -> web.StreamResponse:
+        self.calls.append(("pause_download", None))
+        return web.json_response({"operation": "pause_download"})
+
+    async def resume_download(self, request: web.Request) -> web.StreamResponse:
+        self.calls.append(("resume_download", None))
+        return web.json_response({"operation": "resume_download"})
+
+    async def start_force_download(self, request: web.Request) -> web.StreamResponse:
+        payload = await request.json()
+        self.calls.append(("start_force_download", payload))
+        return web.json_response({"operation": "start_force_download", "payload": payload})
+
+
+class StubExampleImagesProcessor:
+    def __init__(self) -> None:
+        self.calls: List[Tuple[str, Any]] = []
+
+    async def import_images(self, request: web.Request) -> web.StreamResponse:
+        payload = await request.json()
+        self.calls.append(("import_images", payload))
+        return web.json_response({"operation": "import_images", "payload": payload})
+
+    async def delete_custom_image(self, request: web.Request) -> web.StreamResponse:
+        payload = await request.json()
+        self.calls.append(("delete_custom_image", payload))
+        return web.json_response({"operation": "delete_custom_image", "payload": payload})
+
+
+class StubExampleImagesFileManager:
+    def __init__(self) -> None:
+        self.calls: List[Tuple[str, Any]] = []
+
+    async def open_folder(self, request: web.Request) -> web.StreamResponse:
+        payload = await request.json()
+        self.calls.append(("open_folder", payload))
+        return web.json_response({"operation": "open_folder", "payload": payload})
+
+    async def get_files(self, request: web.Request) -> web.StreamResponse:
+        self.calls.append(("get_files", dict(request.query)))
+        return web.json_response({"operation": "get_files", "query": dict(request.query)})
+
+    async def has_images(self, request: web.Request) -> web.StreamResponse:
+        self.calls.append(("has_images", dict(request.query)))
+        return web.json_response({"operation": "has_images", "query": dict(request.query)})
+
+
+@asynccontextmanager
+async def example_images_app() -> ExampleImagesHarness:
+    """Yield an ExampleImagesRoutes app wired with stubbed collaborators."""
+
+    download_manager = StubDownloadManager()
+    processor = StubExampleImagesProcessor()
+    file_manager = StubExampleImagesFileManager()
+
+    controller = ExampleImagesRoutes(
+        download_manager=download_manager,
+        processor=processor,
+        file_manager=file_manager,
+    )
+
+    app = web.Application()
+    controller.register(app)
+
+    server = TestServer(app)
+    client = TestClient(server)
+    await client.start_server()
+
+    try:
+        yield ExampleImagesHarness(
+            client=client,
+            download_manager=download_manager,
+            processor=processor,
+            file_manager=file_manager,
+            controller=controller,
+        )
+    finally:
+        await client.close()
+
+
+async def test_setup_routes_registers_all_definitions():
+    async with example_images_app() as harness:
+        registered = {
+            (route.method, route.resource.canonical)
+            for route in harness.client.app.router.routes()
+            if route.resource.canonical
+        }
+
+        expected = {(definition.method, definition.path) for definition in ROUTE_DEFINITIONS}
+
+        assert expected <= registered
+
+
+@pytest.mark.parametrize(
+    "endpoint, payload",
+    [
+        ("/api/lm/download-example-images", {"model_types": ["lora"], "optimize": False}),
+        ("/api/lm/force-download-example-images", {"model_hashes": ["abc123"]}),
+    ],
+)
+async def test_download_routes_delegate_to_manager(endpoint, payload):
+    async with example_images_app() as harness:
+        response = await harness.client.post(endpoint, json=payload)
+        body = await response.json()
+
+        assert response.status == 200
+        assert body["payload"] == payload
+        assert body["operation"].startswith("start")
+
+        expected_call = body["operation"], payload
+        assert expected_call in harness.download_manager.calls
+
+
+async def test_status_route_returns_manager_payload():
+    async with example_images_app() as harness:
+        response = await harness.client.get(
+            "/api/lm/example-images-status", params={"detail": "true"}
+        )
+        body = await response.json()
+
+        assert response.status == 200
+        assert body == {"operation": "get_status"}
+        assert harness.download_manager.calls == [("get_status", {"detail": "true"})]
+
+
+async def test_pause_and_resume_routes_delegate():
+    async with example_images_app() as harness:
+        pause_response = await harness.client.post("/api/lm/pause-example-images")
+        resume_response = await harness.client.post("/api/lm/resume-example-images")
+
+        assert pause_response.status == 200
+        assert await pause_response.json() == {"operation": "pause_download"}
+        assert resume_response.status == 200
+        assert await resume_response.json() == {"operation": "resume_download"}
+
+        assert harness.download_manager.calls[-2:] == [
+            ("pause_download", None),
+            ("resume_download", None),
+        ]
+
+
+async def test_import_route_delegates_to_processor():
+    payload = {"model_hash": "abc123", "files": ["/path/image.png"]}
+    async with example_images_app() as harness:
+        response = await harness.client.post(
+            "/api/lm/import-example-images", json=payload
+        )
+        body = await response.json()
+
+        assert response.status == 200
+        assert body == {"operation": "import_images", "payload": payload}
+        assert harness.processor.calls == [("import_images", payload)]
+
+
+async def test_delete_route_delegates_to_processor():
+    payload = {"model_hash": "abc123", "short_id": "xyz"}
+    async with example_images_app() as harness:
+        response = await harness.client.post(
+            "/api/lm/delete-example-image", json=payload
+        )
+        body = await response.json()
+
+        assert response.status == 200
+        assert body == {"operation": "delete_custom_image", "payload": payload}
+        assert harness.processor.calls == [("delete_custom_image", payload)]
+
+
+async def test_file_routes_delegate_to_file_manager():
+    open_payload = {"model_hash": "abc123"}
+    files_params = {"model_hash": "def456"}
+
+    async with example_images_app() as harness:
+        open_response = await harness.client.post(
+            "/api/lm/open-example-images-folder", json=open_payload
+        )
+        files_response = await harness.client.get(
+            "/api/lm/example-image-files", params=files_params
+        )
+        has_response = await harness.client.get(
+            "/api/lm/has-example-images", params=files_params
+        )
+
+        assert open_response.status == 200
+        assert files_response.status == 200
+        assert has_response.status == 200
+
+        assert await open_response.json() == {"operation": "open_folder", "payload": open_payload}
+        assert await files_response.json() == {
+            "operation": "get_files",
+            "query": files_params,
+        }
+        assert await has_response.json() == {
+            "operation": "has_images",
+            "query": files_params,
+        }
+
+        assert harness.file_manager.calls == [
+            ("open_folder", open_payload),
+            ("get_files", files_params),
+            ("has_images", files_params),
+        ]
+
+
+@pytest.mark.asyncio
+async def test_download_handler_methods_delegate() -> None:
+    class Recorder:
+        def __init__(self) -> None:
+            self.calls: List[Tuple[str, Any]] = []
+
+        async def start_download(self, request) -> str:
+            self.calls.append(("start_download", request))
+            return "download"
+
+        async def get_status(self, request) -> str:
+            self.calls.append(("get_status", request))
+            return "status"
+
+        async def pause_download(self, request) -> str:
+            self.calls.append(("pause_download", request))
+            return "pause"
+
+        async def resume_download(self, request) -> str:
+            self.calls.append(("resume_download", request))
+            return "resume"
+
+        async def start_force_download(self, request) -> str:
+            self.calls.append(("start_force_download", request))
+            return "force"
+
+    recorder = Recorder()
+    handler = ExampleImagesDownloadHandler(recorder)
+    request = object()
+
+    assert await handler.download_example_images(request) == "download"
+    assert await handler.get_example_images_status(request) == "status"
+    assert await handler.pause_example_images(request) == "pause"
+    assert await handler.resume_example_images(request) == "resume"
+    assert await handler.force_download_example_images(request) == "force"
+
+    expected = [
+        ("start_download", request),
+        ("get_status", request),
+        ("pause_download", request),
+        ("resume_download", request),
+        ("start_force_download", request),
+    ]
+    assert recorder.calls == expected
+
+
+@pytest.mark.asyncio
+async def test_management_handler_methods_delegate() -> None:
+    class Recorder:
+        def __init__(self) -> None:
+            self.calls: List[Tuple[str, Any]] = []
+
+        async def import_images(self, request) -> str:
+            self.calls.append(("import_images", request))
+            return "import"
+
+        async def delete_custom_image(self, request) -> str:
+            self.calls.append(("delete_custom_image", request))
+            return "delete"
+
+    recorder = Recorder()
+    handler = ExampleImagesManagementHandler(recorder)
+    request = object()
+
+    assert await handler.import_example_images(request) == "import"
+    assert await handler.delete_example_image(request) == "delete"
+    assert recorder.calls == [
+        ("import_images", request),
+        ("delete_custom_image", request),
+    ]
+
+
+@pytest.mark.asyncio
+async def test_file_handler_methods_delegate() -> None:
+    class Recorder:
+        def __init__(self) -> None:
+            self.calls: List[Tuple[str, Any]] = []
+
+        async def open_folder(self, request) -> str:
+            self.calls.append(("open_folder", request))
+            return "open"
+
+        async def get_files(self, request) -> str:
+            self.calls.append(("get_files", request))
+            return "files"
+
+        async def has_images(self, request) -> str:
+            self.calls.append(("has_images", request))
+            return "has"
+
+    recorder = Recorder()
+    handler = ExampleImagesFileHandler(recorder)
+    request = object()
+
+    assert await handler.open_example_images_folder(request) == "open"
+    assert await handler.get_example_image_files(request) == "files"
+    assert await handler.has_example_images(request) == "has"
+    assert recorder.calls == [
+        ("open_folder", request),
+        ("get_files", request),
+        ("has_images", request),
+    ]
+
+
+def test_handler_set_route_mapping_includes_all_handlers() -> None:
+    download = ExampleImagesDownloadHandler(object())
+    management = ExampleImagesManagementHandler(object())
+    files = ExampleImagesFileHandler(object())
+    handler_set = ExampleImagesHandlerSet(
+        download=download,
+        management=management,
+        files=files,
+    )
+
+    mapping = handler_set.to_route_mapping()
+
+    expected_keys = {
+        "download_example_images",
+        "get_example_images_status",
+        "pause_example_images",
+        "resume_example_images",
+        "force_download_example_images",
+        "import_example_images",
+        "delete_example_image",
+        "open_example_images_folder",
+        "get_example_image_files",
+        "has_example_images",
+    }
+
+    assert mapping.keys() == expected_keys
+    for key in expected_keys:
+        assert callable(mapping[key])