1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
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"]
|
