mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-03-24 06:32:12 -03:00
235 lines
8.9 KiB
Python
235 lines
8.9 KiB
Python
from contextlib import asynccontextmanager
|
|
from dataclasses import dataclass
|
|
from typing import Any, List, Set, Tuple
|
|
|
|
from aiohttp import web
|
|
from aiohttp.test_utils import TestClient, TestServer
|
|
import pytest
|
|
|
|
from py.routes import example_images_routes
|
|
from py.routes.example_images_routes import ExampleImagesRoutes
|
|
from py.routes.example_images_route_registrar import ROUTE_DEFINITIONS
|
|
|
|
|
|
@dataclass
|
|
class ExampleImagesHarness:
|
|
"""Container exposing the aiohttp client and stubbed collaborators."""
|
|
|
|
client: TestClient
|
|
download_manager: Any
|
|
processor: Any
|
|
file_manager: Any
|
|
|
|
|
|
@asynccontextmanager
|
|
async def example_images_app(monkeypatch: pytest.MonkeyPatch) -> ExampleImagesHarness:
|
|
"""Yield an ExampleImagesRoutes app wired with stubbed collaborators."""
|
|
|
|
class StubDownloadManager:
|
|
calls: List[Tuple[str, Any]] = []
|
|
|
|
@staticmethod
|
|
async def start_download(request):
|
|
payload = await request.json()
|
|
StubDownloadManager.calls.append(("start_download", payload))
|
|
return web.json_response({"operation": "start_download", "payload": payload})
|
|
|
|
@staticmethod
|
|
async def get_status(request):
|
|
StubDownloadManager.calls.append(("get_status", dict(request.query)))
|
|
return web.json_response({"operation": "get_status"})
|
|
|
|
@staticmethod
|
|
async def pause_download(request):
|
|
StubDownloadManager.calls.append(("pause_download", None))
|
|
return web.json_response({"operation": "pause_download"})
|
|
|
|
@staticmethod
|
|
async def resume_download(request):
|
|
StubDownloadManager.calls.append(("resume_download", None))
|
|
return web.json_response({"operation": "resume_download"})
|
|
|
|
@staticmethod
|
|
async def start_force_download(request):
|
|
payload = await request.json()
|
|
StubDownloadManager.calls.append(("start_force_download", payload))
|
|
return web.json_response({"operation": "start_force_download", "payload": payload})
|
|
|
|
class StubExampleImagesProcessor:
|
|
calls: List[Tuple[str, Any]] = []
|
|
|
|
@staticmethod
|
|
async def import_images(request):
|
|
payload = await request.json()
|
|
StubExampleImagesProcessor.calls.append(("import_images", payload))
|
|
return web.json_response({"operation": "import_images", "payload": payload})
|
|
|
|
@staticmethod
|
|
async def delete_custom_image(request):
|
|
payload = await request.json()
|
|
StubExampleImagesProcessor.calls.append(("delete_custom_image", payload))
|
|
return web.json_response({"operation": "delete_custom_image", "payload": payload})
|
|
|
|
class StubExampleImagesFileManager:
|
|
calls: List[Tuple[str, Any]] = []
|
|
|
|
@staticmethod
|
|
async def open_folder(request):
|
|
payload = await request.json()
|
|
StubExampleImagesFileManager.calls.append(("open_folder", payload))
|
|
return web.json_response({"operation": "open_folder", "payload": payload})
|
|
|
|
@staticmethod
|
|
async def get_files(request):
|
|
StubExampleImagesFileManager.calls.append(("get_files", dict(request.query)))
|
|
return web.json_response({"operation": "get_files", "query": dict(request.query)})
|
|
|
|
@staticmethod
|
|
async def has_images(request):
|
|
StubExampleImagesFileManager.calls.append(("has_images", dict(request.query)))
|
|
return web.json_response({"operation": "has_images", "query": dict(request.query)})
|
|
|
|
monkeypatch.setattr(example_images_routes, "DownloadManager", StubDownloadManager)
|
|
monkeypatch.setattr(example_images_routes, "ExampleImagesProcessor", StubExampleImagesProcessor)
|
|
monkeypatch.setattr(example_images_routes, "ExampleImagesFileManager", StubExampleImagesFileManager)
|
|
|
|
app = web.Application()
|
|
ExampleImagesRoutes.setup_routes(app)
|
|
|
|
server = TestServer(app)
|
|
client = TestClient(server)
|
|
await client.start_server()
|
|
|
|
try:
|
|
yield ExampleImagesHarness(
|
|
client=client,
|
|
download_manager=StubDownloadManager,
|
|
processor=StubExampleImagesProcessor,
|
|
file_manager=StubExampleImagesFileManager,
|
|
)
|
|
finally:
|
|
await client.close()
|
|
|
|
|
|
async def test_setup_routes_registers_all_definitions(monkeypatch: pytest.MonkeyPatch):
|
|
async with example_images_app(monkeypatch) as harness:
|
|
registered: Set[tuple[str, str]] = {
|
|
(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, monkeypatch: pytest.MonkeyPatch):
|
|
async with example_images_app(monkeypatch) 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(monkeypatch: pytest.MonkeyPatch):
|
|
async with example_images_app(monkeypatch) 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(monkeypatch: pytest.MonkeyPatch):
|
|
async with example_images_app(monkeypatch) 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(monkeypatch: pytest.MonkeyPatch):
|
|
payload = {"model_hash": "abc123", "files": ["/path/image.png"]}
|
|
async with example_images_app(monkeypatch) 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(monkeypatch: pytest.MonkeyPatch):
|
|
payload = {"model_hash": "abc123", "short_id": "xyz"}
|
|
async with example_images_app(monkeypatch) 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(monkeypatch: pytest.MonkeyPatch):
|
|
open_payload = {"model_hash": "abc123"}
|
|
files_params = {"model_hash": "def456"}
|
|
|
|
async with example_images_app(monkeypatch) 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),
|
|
]
|