aboutsummaryrefslogtreecommitdiff
diff options
context:
space:
mode:
Diffstat
-rw-r--r--.gitignore7
-rw-r--r--README.md17
-rw-r--r--comfier.py20
-rw-r--r--devkit/__init__.py5
-rw-r--r--devkit/fake_open_webui.py96
-rw-r--r--fixtures/nodes.json7
-rw-r--r--fixtures/workflow.json (renamed from workflow.json)0
-rw-r--r--harness.py166
-rw-r--r--pyproject.toml31
-rw-r--r--requirements.txt2
-rw-r--r--scripts/check_stubs.py172
-rw-r--r--stubs/fastapi/__init__.pyi1
-rw-r--r--stubs/open_webui/__init__.pyi0
-rw-r--r--stubs/open_webui/models/__init__.pyi0
-rw-r--r--stubs/open_webui/models/config.pyi5
-rw-r--r--stubs/open_webui/models/users.pyi11
-rw-r--r--stubs/open_webui/routers/__init__.pyi0
-rw-r--r--stubs/open_webui/routers/images.pyi14
-rw-r--r--tests/__init__.py0
-rw-r--r--tests/conftest.py144
-rw-r--r--tests/test_generate.py93
-rw-r--r--tests/test_helpers.py59
22 files changed, 719 insertions, 131 deletions
diff --git a/.gitignore b/.gitignore
index ba0430d..30ed96d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1 +1,6 @@
-__pycache__/ \ No newline at end of file
+__pycache__/
+test_out/
+.pytest_cache/
+.ruff_cache/
+.venv/
+*.egg-info/
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..ed51e3f
--- /dev/null
+++ b/README.md
@@ -0,0 +1,17 @@
+# Comfier
+
+Development repo for the Comfier Open WebUI tool (`comfier.py`). To deploy, paste `comfier.py` into Open WebUI.
+
+## Setup
+ source .venv/bin/activate
+ pip install -e . --group dev
+
+## Workflow
+- Run against a real ComfyUI: `python harness.py --url http://host:8188` (see `--help`; `--cancel-after 2` simulates Stop).
+ Node mappings come from `fixtures/nodes.json` (mirrors `COMFYUI_WORKFLOW_NODES`), workflow from `fixtures/workflow.json`.
+- Offline tests (in-process fake ComfyUI): `pytest`
+- Type check: `pyright` (Open WebUI/FastAPI types come from `stubs/`)
+- Verify stubs against upstream: `python scripts/check_stubs.py` (needs network; checks against the tag in `required_open_webui_version`, or pass `--ref main`). Run it before bumping that version.
+- Lint/format: `ruff check . && ruff format .`
+
+`devkit/` provides the fake `open_webui.*` / `fastapi` modules used by the harness and tests.
diff --git a/comfier.py b/comfier.py
index a97fd0d..33277fb 100644
--- a/comfier.py
+++ b/comfier.py
@@ -8,7 +8,6 @@ version: 0.1.0
license: 3-clause-BSD
"""
-import aiohttp
import asyncio
import copy
import inspect
@@ -17,14 +16,14 @@ import mimetypes
import random
import time
import uuid
-
-from fastapi import Request
-from pydantic import BaseModel, Field
from typing import Annotated, Any, Awaitable, Callable, Optional
+import aiohttp
+from fastapi import Request
from open_webui.models.config import Config
from open_webui.models.users import Users
from open_webui.routers.images import get_image_config, upload_image
+from pydantic import BaseModel, Field
EventEmitter = Callable[[dict[str, Any]], Awaitable[None]]
EventCall = Callable[[dict[str, Any]], Awaitable[Any]]
@@ -286,7 +285,7 @@ class Tools:
status_info = (h_data[prompt_id] or {}).get("status") or {}
if status_info.get("status_str") == "error":
raise RuntimeError(
- "The job was terminated win an error or interruption on the backend."
+ "The job was terminated with an error or interruption on the backend."
)
return True
@@ -378,7 +377,7 @@ class Tools:
if not is_alive:
raise ConnectionError(
f"ComfyUI backend lost prompt {prompt_id} (backend may have restarted or crashed)."
- )
+ ) from None
continue
if msg.type != aiohttp.WSMsgType.TEXT:
@@ -457,7 +456,7 @@ class Tools:
if not output_images:
raise RuntimeError(
- "Generation finished without producing any output. The Job may have beend cancelled or removed from the queue."
+ "Generation finished without producing any output. The Job may have been cancelled or removed from the queue."
)
await self._emit_status(
@@ -486,7 +485,7 @@ class Tools:
or "image/png"
)
- uploaded = upload_image(
+ uploaded: Any = upload_image(
request=request,
image_data=image_bytes,
content_type=content_type,
@@ -502,10 +501,10 @@ class Tools:
if isinstance(uploaded, str):
url = uploaded
- elif hasattr(uploaded, "id"):
- url = f"/api/v1/files/{uploaded.id}/content"
elif isinstance(uploaded, dict) and "id" in uploaded:
url = f"/api/v1/files/{uploaded['id']}/content"
+ elif hasattr(uploaded, "id"):
+ url = f"/api/v1/files/{uploaded.id}/content" # pyright: ignore[reportAttributeAccessIssue]
else:
raise RuntimeError(
f"Unexpected return type from upload_image: {type(uploaded)!r}"
@@ -699,4 +698,3 @@ class Tools:
raw_steps = await Config.get("image_generation.steps", 25)
return int(raw_steps)
-
diff --git a/devkit/__init__.py b/devkit/__init__.py
new file mode 100644
index 0000000..ad4cfbc
--- /dev/null
+++ b/devkit/__init__.py
@@ -0,0 +1,5 @@
+"""Development kit: fake Open WebUI environment for running Comfier outside Open WebUI."""
+
+from devkit.fake_open_webui import FakeEnv, install
+
+__all__ = ["FakeEnv", "install"]
diff --git a/devkit/fake_open_webui.py b/devkit/fake_open_webui.py
new file mode 100644
index 0000000..eb0a1b0
--- /dev/null
+++ b/devkit/fake_open_webui.py
@@ -0,0 +1,96 @@
+"""Fake `fastapi` / `open_webui.*` modules, installed into `sys.modules`."""
+
+import pathlib
+import sys
+import time
+import types
+from dataclasses import dataclass, field
+from types import SimpleNamespace
+from typing import Any
+
+
+@dataclass
+class FakeEnv:
+ """Stand-in for the Open WebUI configuration the tool reads."""
+
+ comfyui_url: str
+ workflow: str
+ api_key: str = ""
+ nodes: list[dict[str, Any]] = field(default_factory=list)
+ image_size: str = "1024x1024"
+ image_steps: int = 25
+ output_dir: pathlib.Path = pathlib.Path("test_out")
+ stored: list[pathlib.Path] = field(default_factory=list)
+
+
+def _module(name: str, **attrs: Any) -> types.ModuleType:
+ mod = types.ModuleType(name)
+ for key, value in attrs.items():
+ setattr(mod, key, value)
+ return mod
+
+
+def install(env: FakeEnv) -> None:
+ """Register fake modules. Call before importing `comfier`."""
+
+ class Request: ...
+
+ class Config:
+ @staticmethod
+ async def get(key: str, default: Any = None) -> Any:
+ return default
+
+ class UsersTable:
+ async def get_user_by_id(self, id: str, db: Any = None) -> Any:
+ return SimpleNamespace(id=id, name="dev")
+
+ Users = UsersTable()
+
+ async def get_image_config() -> Any:
+ return SimpleNamespace(
+ COMFYUI_BASE_URL=env.comfyui_url,
+ COMFYUI_API_KEY=env.api_key,
+ COMFYUI_WORKFLOW=env.workflow,
+ COMFYUI_WORKFLOW_NODES=env.nodes,
+ IMAGE_SIZE=env.image_size,
+ IMAGE_STEPS=env.image_steps,
+ )
+
+ async def upload_image(
+ request: Any,
+ image_data: bytes,
+ content_type: str,
+ metadata: dict[str, Any],
+ user: Any,
+ ) -> Any:
+ env.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')}_{time.monotonic_ns()}"
+ out_file = env.output_dir / f"{file_id}{ext}"
+ out_file.write_bytes(image_data)
+ env.stored.append(out_file)
+ print(f" [Storage] Saved {len(image_data)} bytes -> {out_file}")
+ # Open WebUI 0.12 returns (file_item, {id, url, name, content_type})
+ return SimpleNamespace(id=file_id), {
+ "id": file_id,
+ "url": f"/api/v1/files/{file_id}/content",
+ }
+
+ modules = {
+ "fastapi": _module("fastapi", Request=Request),
+ "open_webui": _module("open_webui"),
+ "open_webui.models": _module("open_webui.models"),
+ "open_webui.routers": _module("open_webui.routers"),
+ "open_webui.models.config": _module("open_webui.models.config", Config=Config),
+ "open_webui.models.users": _module("open_webui.models.users", Users=Users),
+ "open_webui.routers.images": _module(
+ "open_webui.routers.images",
+ get_image_config=get_image_config,
+ upload_image=upload_image,
+ ),
+ }
+ sys.modules.pop("comfier", None)
+ sys.modules.update(modules)
+
+
+__all__ = ["FakeEnv", "install"]
diff --git a/fixtures/nodes.json b/fixtures/nodes.json
new file mode 100644
index 0000000..0e03ae6
--- /dev/null
+++ b/fixtures/nodes.json
@@ -0,0 +1,7 @@
+[
+ {"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"]}
+]
diff --git a/workflow.json b/fixtures/workflow.json
index 40d24a8..40d24a8 100644
--- a/workflow.json
+++ b/fixtures/workflow.json
diff --git a/harness.py b/harness.py
index f588328..450df24 100644
--- a/harness.py
+++ b/harness.py
@@ -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,
diff --git a/pyproject.toml b/pyproject.toml
new file mode 100644
index 0000000..295ae49
--- /dev/null
+++ b/pyproject.toml
@@ -0,0 +1,31 @@
+[project]
+name = "comfier"
+version = "0.1.0"
+description = "Development repository for the Comfier Open WebUI tool"
+requires-python = ">=3.12"
+dependencies = ["aiohttp", "pydantic"]
+
+[dependency-groups]
+dev = ["pyright", "ruff", "pytest", "pytest-asyncio"]
+
+[tool.setuptools]
+py-modules = ["comfier", "harness"]
+packages = ["devkit"]
+
+[tool.pyright]
+include = ["comfier.py", "harness.py", "devkit", "tests", "scripts"]
+stubPath = "stubs"
+typeCheckingMode = "standard"
+reportMissingModuleSource = false
+pythonVersion = "3.14"
+
+[tool.ruff]
+extend-exclude = ["stubs"]
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "B"]
+ignore = ["E501"]
+
+[tool.pytest.ini_options]
+asyncio_mode = "auto"
+testpaths = ["tests"]
diff --git a/requirements.txt b/requirements.txt
deleted file mode 100644
index 7925b0b..0000000
--- a/requirements.txt
+++ /dev/null
@@ -1,2 +0,0 @@
-aiohttp
-pydantic \ No newline at end of file
diff --git a/scripts/check_stubs.py b/scripts/check_stubs.py
new file mode 100644
index 0000000..5d5d7e4
--- /dev/null
+++ b/scripts/check_stubs.py
@@ -0,0 +1,172 @@
+#!/usr/bin/env python3
+"""Check that stubs/ match the real Open WebUI signatures.
+
+Fetches the relevant Open WebUI source files from GitHub at a given ref
+(default: the tag from `required_open_webui_version` in comfier.py), parses
+them with `ast` (nothing is imported or installed) and compares function
+signatures with the .pyi stubs: async-ness, parameter names/kinds and whether
+each parameter has a default. Annotations are not compared.
+"""
+
+import argparse
+import ast
+import pathlib
+import re
+import sys
+import urllib.request
+
+ROOT = pathlib.Path(__file__).parent.parent
+RAW = "https://raw.githubusercontent.com/open-webui/open-webui/{ref}/backend/open_webui/{path}.py"
+
+# (stub file, upstream module path, owner class or None, function names)
+CHECKS: list[tuple[str, str, str | None, list[str]]] = [
+ ("models/config", "models/config", "Config", ["get"]),
+ ("models/users", "models/users", "UsersTable", ["get_user_by_id"]),
+ ("routers/images", "routers/images", None, ["get_image_config", "upload_image"]),
+]
+
+Fn = ast.FunctionDef | ast.AsyncFunctionDef
+
+
+def default_ref() -> str:
+ header = (ROOT / "comfier.py").read_text(encoding="utf-8")
+ match = re.search(r"^required_open_webui_version:\s*(\S+)", header, re.MULTILINE)
+ if not match:
+ sys.exit("Could not find required_open_webui_version in comfier.py")
+ return "v" + match.group(1).lstrip("v")
+
+
+def fetch(ref: str, path: str) -> str:
+ url = RAW.format(ref=ref, path=path)
+ with urllib.request.urlopen(url) as resp:
+ return resp.read().decode("utf-8")
+
+
+def find_fn(tree: ast.Module, owner: str | None, name: str) -> Fn | None:
+ body: list[ast.stmt] = tree.body
+ if owner is not None:
+ cls = next(
+ (n for n in body if isinstance(n, ast.ClassDef) and n.name == owner), None
+ )
+ if cls is None:
+ return None
+ body = cls.body
+ return next(
+ (
+ n
+ for n in body
+ if isinstance(n, ast.FunctionDef | ast.AsyncFunctionDef) and n.name == name
+ ),
+ None,
+ )
+
+
+def signature(fn: Fn, drop_self: bool) -> tuple[bool, list[tuple[str, str, bool]]]:
+ """(is_async, [(name, kind, has_default)])"""
+ a = fn.args
+ params: list[tuple[str, str, bool]] = []
+ positional = [*a.posonlyargs, *a.args]
+ n_defaults = len(a.defaults)
+ for i, arg in enumerate(positional):
+ has_default = i >= len(positional) - n_defaults
+ kind = "posonly" if arg in a.posonlyargs else "pos"
+ params.append((arg.arg, kind, has_default))
+ if a.vararg:
+ params.append((a.vararg.arg, "vararg", False))
+ for arg, default in zip(a.kwonlyargs, a.kw_defaults, strict=True):
+ params.append((arg.arg, "kwonly", default is not None))
+ if a.kwarg:
+ params.append((a.kwarg.arg, "kwarg", False))
+ if drop_self and params and params[0][0] in ("self", "cls"):
+ params = params[1:]
+ return isinstance(fn, ast.AsyncFunctionDef), params
+
+
+def check_users_singleton(upstream: ast.Module, stub: ast.Module) -> list[str]:
+ """`Users` must be an instance of UsersTable in both."""
+ problems: list[str] = []
+ real = any(
+ isinstance(n, ast.Assign)
+ and any(isinstance(t, ast.Name) and t.id == "Users" for t in n.targets)
+ and isinstance(n.value, ast.Call)
+ and isinstance(n.value.func, ast.Name)
+ and n.value.func.id == "UsersTable"
+ for n in upstream.body
+ )
+ stubbed = any(
+ isinstance(n, ast.AnnAssign)
+ and isinstance(n.target, ast.Name)
+ and n.target.id == "Users"
+ and isinstance(n.annotation, ast.Name)
+ and n.annotation.id == "UsersTable"
+ for n in stub.body
+ )
+ if not real:
+ problems.append("models/users: upstream no longer has `Users = UsersTable()`")
+ if not stubbed:
+ problems.append("models/users: stub lacks `Users: UsersTable`")
+ return problems
+
+
+def is_static(fn: Fn) -> bool:
+ return any(
+ isinstance(d, ast.Name) and d.id == "staticmethod" for d in fn.decorator_list
+ )
+
+
+def main() -> int:
+ parser = argparse.ArgumentParser(description=(__doc__ or "").splitlines()[0])
+ parser.add_argument(
+ "--ref", default=None, help="git ref/tag (default: from comfier.py)"
+ )
+ args = parser.parse_args()
+ ref = args.ref or default_ref()
+ print(f"Checking stubs against open-webui @ {ref}")
+
+ problems: list[str] = []
+ for stub_path, upstream_path, owner, names in CHECKS:
+ stub_file = ROOT / "stubs" / "open_webui" / f"{stub_path}.pyi"
+ stub_tree = ast.parse(stub_file.read_text(encoding="utf-8"))
+ try:
+ upstream_tree = ast.parse(fetch(ref, upstream_path))
+ except Exception as err:
+ problems.append(f"{upstream_path}: could not fetch/parse upstream ({err})")
+ continue
+
+ for name in names:
+ label = f"{upstream_path}:{owner + '.' if owner else ''}{name}"
+ real_fn = find_fn(upstream_tree, owner, name)
+ stub_fn = find_fn(stub_tree, owner, name)
+ if real_fn is None:
+ problems.append(f"{label}: missing upstream")
+ continue
+ if stub_fn is None:
+ problems.append(f"{label}: missing in stub")
+ continue
+ real_sig = signature(
+ real_fn, drop_self=owner is not None and not is_static(real_fn)
+ )
+ stub_sig = signature(
+ stub_fn, drop_self=owner is not None and not is_static(stub_fn)
+ )
+ if real_sig != stub_sig:
+ problems.append(
+ f"{label}:\n upstream: {real_sig}\n stub: {stub_sig}"
+ )
+ else:
+ print(f" ok {label}")
+
+ if owner == "UsersTable":
+ problems += check_users_singleton(upstream_tree, stub_tree)
+
+ if problems:
+ print("\nMismatches:")
+ for p in problems:
+ print(f" - {p}")
+ return 1
+ print("All stubs match.")
+ return 0
+
+
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/stubs/fastapi/__init__.pyi b/stubs/fastapi/__init__.pyi
new file mode 100644
index 0000000..2eda7fc
--- /dev/null
+++ b/stubs/fastapi/__init__.pyi
@@ -0,0 +1 @@
+class Request: ...
diff --git a/stubs/open_webui/__init__.pyi b/stubs/open_webui/__init__.pyi
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/stubs/open_webui/__init__.pyi
diff --git a/stubs/open_webui/models/__init__.pyi b/stubs/open_webui/models/__init__.pyi
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/stubs/open_webui/models/__init__.pyi
diff --git a/stubs/open_webui/models/config.pyi b/stubs/open_webui/models/config.pyi
new file mode 100644
index 0000000..ae951b2
--- /dev/null
+++ b/stubs/open_webui/models/config.pyi
@@ -0,0 +1,5 @@
+from typing import Any
+
+class Config:
+ @staticmethod
+ async def get(key: str, default: Any = None) -> Any: ...
diff --git a/stubs/open_webui/models/users.pyi b/stubs/open_webui/models/users.pyi
new file mode 100644
index 0000000..79d409e
--- /dev/null
+++ b/stubs/open_webui/models/users.pyi
@@ -0,0 +1,11 @@
+from typing import Any
+
+class UserModel:
+ id: str
+ name: str
+ def __getattr__(self, name: str) -> Any: ...
+
+class UsersTable:
+ async def get_user_by_id(self, id: str, db: Any = None) -> UserModel | None: ...
+
+Users: UsersTable
diff --git a/stubs/open_webui/routers/__init__.pyi b/stubs/open_webui/routers/__init__.pyi
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/stubs/open_webui/routers/__init__.pyi
diff --git a/stubs/open_webui/routers/images.pyi b/stubs/open_webui/routers/images.pyi
new file mode 100644
index 0000000..04c759f
--- /dev/null
+++ b/stubs/open_webui/routers/images.pyi
@@ -0,0 +1,14 @@
+from types import SimpleNamespace
+from typing import Any
+
+from fastapi import Request
+
+async def get_image_config() -> SimpleNamespace: ...
+async def upload_image(
+ request: Request,
+ image_data: bytes,
+ content_type: str,
+ metadata: dict[str, Any],
+ user: Any,
+ db: Any = None,
+) -> tuple[Any, dict[str, Any]]: ...
diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..e69de29
--- /dev/null
+++ b/tests/__init__.py
diff --git a/tests/conftest.py b/tests/conftest.py
new file mode 100644
index 0000000..7370405
--- /dev/null
+++ b/tests/conftest.py
@@ -0,0 +1,144 @@
+import asyncio
+import json
+import pathlib
+from collections.abc import AsyncIterator, Callable
+from typing import Any
+
+import pytest
+import pytest_asyncio
+from aiohttp import WSMsgType, web
+
+from devkit import FakeEnv, install
+
+ROOT = pathlib.Path(__file__).parent.parent
+PNG = bytes.fromhex(
+ "89504e470d0a1a0a0000000d49484452000000010000000108060000001f15c489"
+ "0000000d49444154789c6360000002000001e221bc330000000049454e44ae426082"
+)
+
+
+# Make `import comfier` work at collection time; fixtures re-install per test.
+install(FakeEnv(comfyui_url="", workflow="{}"))
+
+
+class FakeComfy:
+ """Minimal in-process ComfyUI. `script` lists WS messages sent after /prompt."""
+
+ def __init__(self) -> None:
+ self.script: list[dict[str, Any]] = []
+ self.submitted: list[dict[str, Any]] = []
+ self.interrupts: list[Any] = []
+ self.drop_ws = False
+ self.hang = False
+ self.url = ""
+ self._ws: web.WebSocketResponse | None = None
+ self._ws_ready = asyncio.Event()
+
+ def success_script(self, prompt_id: str) -> list[dict[str, Any]]:
+ return [
+ {"type": "executing", "data": {"node": "1", "prompt_id": prompt_id}},
+ {
+ "type": "progress",
+ "data": {"value": 1, "max": 2, "prompt_id": prompt_id},
+ },
+ {
+ "type": "executed",
+ "data": {
+ "prompt_id": prompt_id,
+ "output": {"images": [{"filename": "a.png", "type": "output"}]},
+ },
+ },
+ {"type": "executing", "data": {"node": None, "prompt_id": prompt_id}},
+ ]
+
+ async def ws(self, request: web.Request) -> web.WebSocketResponse:
+ ws = web.WebSocketResponse()
+ await ws.prepare(request)
+ self._ws = ws
+ self._ws_ready.set()
+ async for msg in ws:
+ if msg.type == WSMsgType.ERROR:
+ break
+ return ws
+
+ async def prompt(self, request: web.Request) -> web.Response:
+ body = await request.json()
+ self.submitted.append(body["prompt"])
+ pid = "pid-1"
+ await self._ws_ready.wait()
+ assert self._ws is not None
+ ws = self._ws
+ asyncio.get_running_loop().call_soon(
+ lambda: asyncio.ensure_future(self._play(ws, pid))
+ )
+ return web.json_response({"prompt_id": pid})
+
+ async def _play(self, ws: web.WebSocketResponse, pid: str) -> None:
+ if self.hang:
+ return
+ if self.drop_ws:
+ await ws.close()
+ return
+ for msg in self.script or self.success_script(pid):
+ await ws.send_str(json.dumps(msg))
+
+ async def queue(self, request: web.Request) -> web.Response:
+ if request.method == "POST":
+ return web.json_response({})
+ return web.json_response({"queue_running": [[0, "pid-1"]], "queue_pending": []})
+
+ async def interrupt(self, request: web.Request) -> web.Response:
+ self.interrupts.append(await request.json())
+ return web.json_response({})
+
+ async def history(self, request: web.Request) -> web.Response:
+ return web.json_response({})
+
+ async def view(self, request: web.Request) -> web.Response:
+ return web.Response(body=PNG, content_type="image/png")
+
+
+@pytest_asyncio.fixture
+async def comfy() -> AsyncIterator[FakeComfy]:
+ fake = FakeComfy()
+ app = web.Application()
+ app.add_routes(
+ [
+ web.get("/ws", fake.ws),
+ web.post("/prompt", fake.prompt),
+ web.get("/queue", fake.queue),
+ web.post("/queue", fake.queue),
+ web.post("/interrupt", fake.interrupt),
+ web.get("/history/{id}", fake.history),
+ web.get("/view", fake.view),
+ ]
+ )
+ runner = web.AppRunner(app)
+ await runner.setup()
+ site = web.TCPSite(runner, "127.0.0.1", 0)
+ await site.start()
+ port = site._server.sockets[0].getsockname()[1] # type: ignore[union-attr]
+ fake.url = f"http://127.0.0.1:{port}"
+ yield fake
+ await runner.cleanup()
+
+
+@pytest.fixture
+def make_tools(
+ tmp_path: pathlib.Path, comfy: FakeComfy
+) -> Callable[[], tuple[Any, FakeEnv]]:
+ def factory() -> tuple[Any, FakeEnv]:
+ env = FakeEnv(
+ comfyui_url=comfy.url,
+ workflow=(ROOT / "fixtures" / "workflow.json").read_text(encoding="utf-8"),
+ nodes=json.loads(
+ (ROOT / "fixtures" / "nodes.json").read_text(encoding="utf-8")
+ ),
+ output_dir=tmp_path,
+ )
+ install(env)
+ from comfier import Tools
+
+ return Tools(), env
+
+ return factory
diff --git a/tests/test_generate.py b/tests/test_generate.py
new file mode 100644
index 0000000..37dda32
--- /dev/null
+++ b/tests/test_generate.py
@@ -0,0 +1,93 @@
+import asyncio
+from collections.abc import Callable
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+
+from tests.conftest import FakeComfy
+
+
+async def _run(tools: Any, **kw: Any) -> tuple[dict[str, Any], list[dict[str, Any]]]:
+ events: list[dict[str, Any]] = []
+
+ async def emit(e: dict[str, Any]) -> None:
+ events.append(e)
+
+ result = await tools.comfier_generate_image(
+ positive_prompt="a marble",
+ seed="5",
+ size="64x64",
+ __request__=SimpleNamespace(),
+ __user__={"id": "u"},
+ __event_emitter__=emit,
+ **kw,
+ )
+ return result, events
+
+
+async def test_success(make_tools: Callable[[], Any], comfy: FakeComfy) -> None:
+ tools, env = make_tools()
+ result, events = await _run(tools)
+ assert result["status"] == "success"
+ assert len(result["images"]) == 1 and len(env.stored) == 1
+ assert comfy.submitted[0]["520:517"]["inputs"]["seed"] == 5
+ assert comfy.submitted[0]["520:518"]["inputs"]["value"] == 64
+ assert events[-1]["type"] == "files"
+
+
+async def test_execution_error(make_tools: Callable[[], Any], comfy: FakeComfy) -> None:
+ comfy.script = [
+ {
+ "type": "execution_error",
+ "data": {"prompt_id": "pid-1", "exception_message": "boom"},
+ }
+ ]
+ tools, _ = make_tools()
+ with pytest.raises(RuntimeError, match="boom"):
+ await _run(tools)
+
+
+async def test_ws_drop_retry_declined(
+ make_tools: Callable[[], Any], comfy: FakeComfy
+) -> None:
+ comfy.drop_ws = True
+ tools, _ = make_tools()
+ asked: list[dict[str, Any]] = []
+
+ async def call(e: dict[str, Any]) -> bool:
+ asked.append(e)
+ return False
+
+ with pytest.raises(ConnectionError):
+ await _run(tools, __event_call__=call)
+ assert len(asked) == 1
+
+
+async def test_ws_drop_retry_accepted(
+ make_tools: Callable[[], Any], comfy: FakeComfy
+) -> None:
+ comfy.drop_ws = True
+ tools, _ = make_tools()
+
+ async def call(e: dict[str, Any]) -> bool:
+ comfy.drop_ws = False
+ return True
+
+ result, _ = await _run(tools, __event_call__=call)
+ assert result["status"] == "success"
+
+
+async def test_cancel_interrupts_backend(
+ make_tools: Callable[[], Any], comfy: FakeComfy
+) -> None:
+ comfy.hang = True
+ tools, _ = make_tools()
+ task = asyncio.create_task(_run(tools))
+ while not comfy.submitted:
+ await asyncio.sleep(0.05)
+ await asyncio.sleep(0.1)
+ task.cancel()
+ with pytest.raises(asyncio.CancelledError):
+ await task
+ assert comfy.interrupts == [{"prompt_id": "pid-1"}]
diff --git a/tests/test_helpers.py b/tests/test_helpers.py
new file mode 100644
index 0000000..0c0f875
--- /dev/null
+++ b/tests/test_helpers.py
@@ -0,0 +1,59 @@
+import json
+import pathlib
+
+import pytest
+
+from comfier import Tools
+
+ROOT = pathlib.Path(__file__).parent.parent
+
+
+@pytest.fixture
+def tools(make_tools): # type: ignore[no-untyped-def]
+ return make_tools()[0]
+
+
+def test_normalize_seed(tools: Tools) -> None:
+ assert tools._normalize_seed("42") == 42
+ assert tools._normalize_seed("0x10") == 16
+ assert tools._normalize_seed(7) == 7
+ assert 0 <= tools._normalize_seed(None) < 1 << 63
+ assert 0 <= tools._normalize_seed(-5) < 1 << 63
+ assert tools._normalize_seed((1 << 63) + 3) == 3
+ with pytest.raises(ValueError):
+ tools._normalize_seed("abc")
+
+
+async def test_resolve_dimensions(tools: Tools) -> None:
+ assert await tools._resolve_dimensions("512X768") == (512, 768)
+ assert await tools._resolve_dimensions(None) == (1024, 1024)
+ with pytest.raises(ValueError):
+ await tools._resolve_dimensions("0x10")
+ with pytest.raises(ValueError):
+ await tools._resolve_dimensions("bogus")
+
+
+def test_resolve_node_mappings(tools: Tools) -> None:
+ tools.valves.negative_prompt_node_ids = "1, 2,,"
+ nodes = json.loads((ROOT / "fixtures" / "nodes.json").read_text())
+ m = tools._resolve_node_mappings(nodes)
+ assert m["negative_prompt"]["node_ids"] == ["1", "2"]
+ assert m["positive_prompt"] == {"key": "value", "node_ids": ["520:506"]}
+ assert m["steps"]["node_ids"] == ["459:458"]
+
+
+def test_apply_workflow_parameters(tools: Tools) -> None:
+ wf = json.loads((ROOT / "fixtures" / "workflow.json").read_text())
+ m = tools._resolve_node_mappings(
+ json.loads((ROOT / "fixtures" / "nodes.json").read_text())
+ )
+ applied = tools._apply_workflow_parameters(
+ wf, m, {"seed": 9, "width": 64, "steps": None}
+ )
+ assert wf["520:517"]["inputs"]["seed"] == 9
+ assert wf["520:518"]["inputs"]["value"] == 64
+ assert "steps" not in applied
+ with pytest.raises(KeyError):
+ tools._apply_workflow_parameters(
+ wf, {"seed": {"key": "seed", "node_ids": ["nope"]}}, {"seed": 1}
+ )
generated by cgit v1.2.3 (git 2.25.1) at 2026-11-10 16:28:03 +0000