diff options
Diffstat (limited to 'harness.py')
| -rw-r--r-- | harness.py | 166 |
1 files changed, 49 insertions, 117 deletions
@@ -1,97 +1,19 @@ #!/usr/bin/env python3 +"""Standalone CLI harness: run the Comfier tool against a real ComfyUI.""" + import argparse import asyncio import json +import os import pathlib -import sys -import types from types import SimpleNamespace -from typing import Any - -# --------------------------------------------------------------------------- -# 1. Stub out `fastapi` and `open_webui.*` in sys.modules BEFORE importing tool -# --------------------------------------------------------------------------- -def install_open_webui_stubs( - comfyui_url: str, - api_key: str, - workflow_path: pathlib.Path, - nodes_config: list[dict[str, Any]], - image_size: str, - image_steps: int, - output_dir: pathlib.Path, -) -> None: - # Ensure fastapi.Request exists even in a minimal venv - if "fastapi" not in sys.modules: - fastapi_mod = types.ModuleType("fastapi") - fastapi_mod.Request = object - sys.modules["fastapi"] = fastapi_mod - - workflow_raw = workflow_path.read_text(encoding="utf-8") - - # Mock `open_webui.models.config.Config` - config_mod = types.ModuleType("open_webui.models.config") - - class MockConfig: - @staticmethod - async def get(key: str, default: Any = None) -> Any: - return default - - config_mod.Config = MockConfig - - # Mock `open_webui.models.users.Users` - users_mod = types.ModuleType("open_webui.models.users") - - class MockUsers: - @staticmethod - def get_user_by_id(user_id: str) -> Any: - return SimpleNamespace(id=user_id, name="Felix") - - users_mod.Users = MockUsers - - # Mock `open_webui.routers.images` - images_mod = types.ModuleType("open_webui.routers.images") - - async def mock_get_image_config() -> Any: - return SimpleNamespace( - COMFYUI_BASE_URL=comfyui_url, - COMFYUI_API_KEY=api_key, - COMFYUI_WORKFLOW=workflow_raw, - COMFYUI_WORKFLOW_NODES=nodes_config, - IMAGE_SIZE=image_size, - IMAGE_STEPS=image_steps, - ) +from typing import Any, cast + +from devkit import FakeEnv, install + +FIXTURES = pathlib.Path(__file__).parent / "fixtures" + - def mock_upload_image( - request: Any, - image_data: bytes, - content_type: str, - metadata: dict[str, Any], - user: Any, - ) -> Any: - output_dir.mkdir(parents=True, exist_ok=True) - ext = ".png" if "png" in content_type else ".jpg" - file_id = f"seed_{metadata.get('seed', 'unknown')}_{int(asyncio.get_event_loop().time() * 1000)}" - out_file = output_dir / f"{file_id}{ext}" - out_file.write_bytes(image_data) - print(f" [Storage] Saved {len(image_data)} bytes -> {out_file}") - # Return a mock FileModel with an `.id` attribute just like Open WebUI - return SimpleNamespace(id=file_id) - - images_mod.get_image_config = mock_get_image_config - images_mod.upload_image = mock_upload_image - - # Register package hierarchy in sys.modules - sys.modules["open_webui"] = types.ModuleType("open_webui") - sys.modules["open_webui.models"] = types.ModuleType("open_webui.models") - sys.modules["open_webui.routers"] = types.ModuleType("open_webui.routers") - sys.modules["open_webui.models.config"] = config_mod - sys.modules["open_webui.models.users"] = users_mod - sys.modules["open_webui.routers.images"] = images_mod - - -# --------------------------------------------------------------------------- -# 2. Terminal Event Emitter & Interactive Event Call -# --------------------------------------------------------------------------- async def console_event_emitter(event: dict[str, Any]) -> None: ev_type = event.get("type") data = event.get("data", {}) @@ -105,9 +27,8 @@ async def console_event_emitter(event: dict[str, Any]) -> None: async def console_event_call(event: dict[str, Any]) -> bool: - ev_type = event.get("type") data = event.get("data", {}) - print(f"\n--- MODAL ({ev_type}): {data.get('title')} ---") + print(f"\n--- MODAL ({event.get('type')}): {data.get('title')} ---") print(data.get("message", "")) loop = asyncio.get_running_loop() answer = await loop.run_in_executor( @@ -116,42 +37,53 @@ async def console_event_call(event: dict[str, Any]) -> bool: return answer in ("y", "yes") -# --------------------------------------------------------------------------- -# 3. CLI Runner -# --------------------------------------------------------------------------- async def main() -> None: parser = argparse.ArgumentParser(description="Standalone test harness for Comfier") - parser.add_argument("--url", default="http://10.42.0.146:8188", help="ComfyUI base URL") - parser.add_argument("--workflow", type=pathlib.Path, required=True, help="Path to ComfyUI API workflow JSON") - parser.add_argument("--prompt", default="a tiny blue marble on a white surface", help="Positive prompt") + parser.add_argument( + "--url", + default=os.environ.get("COMFYUI_URL", "http://127.0.0.1:8188"), + help="ComfyUI base URL (default: $COMFYUI_URL or http://127.0.0.1:8188)", + ) + parser.add_argument("--api-key", default=os.environ.get("COMFYUI_API_KEY", "")) + parser.add_argument( + "--workflow", type=pathlib.Path, default=FIXTURES / "workflow.json" + ) + parser.add_argument( + "--nodes", + type=pathlib.Path, + default=FIXTURES / "nodes.json", + help="JSON file mirroring Open WebUI's COMFYUI_WORKFLOW_NODES", + ) + parser.add_argument("--prompt", default="a tiny blue marble on a white surface") parser.add_argument("--neg", default=None, help="Negative prompt") parser.add_argument("--size", default="256x256", help="Image size (WIDTHxHEIGHT)") - parser.add_argument("--seed", default="123", help="Seed") + parser.add_argument("--seed", default="123") parser.add_argument("--steps", type=int, default=30, help="Global IMAGE_STEPS") - parser.add_argument("--neg-nodes", default="459:452", help="Valve: negative_prompt_node_ids") - parser.add_argument("--cancel-after", type=float, default=None, help="Simulate user clicking Stop after N seconds") + parser.add_argument( + "--neg-nodes", default="459:452", help="Valve: negative_prompt_node_ids" + ) + parser.add_argument("--out", type=pathlib.Path, default=pathlib.Path("test_out")) + parser.add_argument( + "--cancel-after", + type=float, + default=None, + help="Simulate user clicking Stop after N seconds", + ) args = parser.parse_args() - # Match your Open WebUI global COMFYUI_WORKFLOW_NODES setup - global_nodes = [ - {"type": "prompt", "key": "value", "node_ids": ["520:506"]}, - {"type": "seed", "key": "seed", "node_ids": ["520:517"]}, - {"type": "width", "key": "value", "node_ids": ["520:518"]}, - {"type": "height", "key": "value", "node_ids": ["520:519"]}, - {"type": "steps", "key": "steps", "node_ids": ["459:458"]}, - ] - - install_open_webui_stubs( - comfyui_url=args.url, - api_key="", - workflow_path=args.workflow, - nodes_config=global_nodes, - image_size="1024x1024", - image_steps=args.steps, - output_dir=pathlib.Path("./test_out"), + install( + FakeEnv( + comfyui_url=args.url, + api_key=args.api_key, + workflow=args.workflow.read_text(encoding="utf-8"), + nodes=json.loads(args.nodes.read_text(encoding="utf-8")), + image_size="1024x1024", + image_steps=args.steps, + output_dir=args.out, + ) ) - # Import your tool AFTER stubs are installed + # Import the tool AFTER the fake Open WebUI modules are installed. from comfier import Tools tools = Tools() @@ -163,7 +95,7 @@ async def main() -> None: negative_prompt=args.neg, size=args.size, seed=args.seed, - __request__=SimpleNamespace(), + __request__=cast(Any, SimpleNamespace()), __user__={"id": "local-dev"}, __event_emitter__=console_event_emitter, __event_call__=console_event_call, |