200 lines
7.0 KiB
Python
200 lines
7.0 KiB
Python
"""Test helpers: a RigShare server without ComfyUI.
|
|
|
|
``rig_server()`` starts a real aiohttp app with RigShare's routes, middleware
|
|
and hub, plus a small stand-in for ComfyUI's /userdata API backed by a temp
|
|
folder, so HTTP and websocket behaviour can be tested end to end.
|
|
"""
|
|
|
|
import contextlib
|
|
import glob
|
|
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
from urllib.parse import quote, unquote
|
|
|
|
from aiohttp import web, ClientSession, CookieJar
|
|
from aiohttp.test_utils import TestServer
|
|
|
|
ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
if ROOT not in sys.path:
|
|
sys.path.insert(0, ROOT)
|
|
|
|
from rigshare_core.store import Store # noqa: E402
|
|
from rigshare_core.hub import Hub # noqa: E402
|
|
from rigshare_core.workspace import Workspace # noqa: E402
|
|
from rigshare_core import routes # noqa: E402
|
|
|
|
PASSWORD = "longpassword"
|
|
|
|
|
|
def enc(path):
|
|
"""Encode a userdata path the way ComfyUI's frontend does."""
|
|
return quote(path, safe="")
|
|
|
|
|
|
class FakePromptServer:
|
|
def __init__(self):
|
|
self.app = web.Application()
|
|
self.sent = []
|
|
|
|
def send_sync(self, event, data, sid=None):
|
|
self.sent.append((event, sid))
|
|
|
|
|
|
def comfy_routes(user_root):
|
|
"""Minimal ComfyUI /userdata routes for the single default user."""
|
|
|
|
def resolve(name):
|
|
if "%" in name:
|
|
name = unquote(name)
|
|
path = os.path.abspath(os.path.join(user_root, name))
|
|
return path if os.path.commonpath((user_root, path)) == user_root else None
|
|
|
|
async def listing(request):
|
|
base = resolve(request.query["dir"])
|
|
out = []
|
|
for full in glob.glob(os.path.join(glob.escape(base), "**", "*"), recursive=True):
|
|
if os.path.isfile(full):
|
|
out.append({"path": os.path.relpath(full, base).replace(os.sep, "/"), "size": 1, "modified": 1})
|
|
return web.json_response(out)
|
|
|
|
async def get(request):
|
|
path = resolve(request.match_info["file"])
|
|
return web.FileResponse(path) if path and os.path.exists(path) else web.Response(status=404)
|
|
|
|
async def post(request):
|
|
path = resolve(request.match_info["file"])
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
with open(path, "wb") as f:
|
|
f.write(await request.read())
|
|
return web.json_response(os.path.relpath(path, user_root))
|
|
|
|
async def delete(request):
|
|
os.remove(resolve(request.match_info["file"]))
|
|
return web.Response(status=204)
|
|
|
|
async def move(request):
|
|
src, dst = resolve(request.match_info["file"]), resolve(request.match_info["dest"])
|
|
os.makedirs(os.path.dirname(dst), exist_ok=True)
|
|
os.rename(src, dst)
|
|
return web.json_response(os.path.relpath(dst, user_root))
|
|
|
|
async def ok(request):
|
|
return web.json_response({"ok": True})
|
|
|
|
table = web.RouteTableDef()
|
|
for prefix in ("", "/api"):
|
|
table.get(prefix + "/userdata")(listing)
|
|
table.get(prefix + "/userdata/{file}")(get)
|
|
table.post(prefix + "/userdata/{file}")(post)
|
|
table.delete(prefix + "/userdata/{file}")(delete)
|
|
table.post(prefix + "/userdata/{file}/move/{dest}")(move)
|
|
table.route("*", prefix + "/prompt")(ok)
|
|
table.route("*", prefix + "/manager/{tail:.*}")(ok)
|
|
table.route("*", prefix + "/customnode/{tail:.*}")(ok)
|
|
return table
|
|
|
|
|
|
class Rig:
|
|
"""A running test server plus convenience helpers."""
|
|
|
|
def __init__(self, server, store, hub, prompt_server, user_root):
|
|
self.server, self.store, self.hub, self.ps, self.user_root = server, store, hub, prompt_server, user_root
|
|
self.base = f"http://{server.host}:{server.port}"
|
|
self.http = None
|
|
self.headers = {}
|
|
|
|
def url(self, path):
|
|
return self.base + path
|
|
|
|
async def setup_admin(self, username="boss"):
|
|
r = await self.http.post(self.url("/rigshare/api/setup"), json={"username": username, "password": PASSWORD})
|
|
assert r.status == 200, await r.text()
|
|
self.headers[username] = {"Authorization": "Bearer " + (await r.json())["token"]}
|
|
return self.headers[username]
|
|
|
|
async def add_user(self, username, **perms):
|
|
admin = next(iter(self.headers.values()))
|
|
r = await self.http.post(self.url("/rigshare/api/users"), headers=admin,
|
|
json={"username": username, "password": PASSWORD, "perms": perms})
|
|
assert r.status == 200, await r.text()
|
|
r = await self.http.post(self.url("/rigshare/api/login"), json={"username": username, "password": PASSWORD})
|
|
self.headers[username] = {"Authorization": "Bearer " + (await r.json())["token"]}
|
|
return self.headers[username]
|
|
|
|
async def ws(self, username):
|
|
session = ClientSession()
|
|
sock = await session.ws_connect(self.url("/rigshare/ws"), headers=self.headers[username])
|
|
await sock.send_json({"type": "hello"})
|
|
await recv_until(sock, "welcome")
|
|
return Socket(session, sock)
|
|
|
|
|
|
class Socket:
|
|
def __init__(self, session, sock):
|
|
self.session, self.sock = session, sock
|
|
|
|
async def send(self, message):
|
|
await self.sock.send_json(message)
|
|
|
|
async def until(self, kind, timeout=2.0):
|
|
return await recv_until(self.sock, kind, timeout)
|
|
|
|
async def drain(self, timeout=0.3):
|
|
return await drain(self.sock, timeout)
|
|
|
|
async def close(self):
|
|
await self.sock.close()
|
|
await self.session.close()
|
|
|
|
|
|
async def recv_until(sock, kind, timeout=2.0):
|
|
import asyncio
|
|
loop = asyncio.get_running_loop()
|
|
end = loop.time() + timeout
|
|
while True:
|
|
msg = await asyncio.wait_for(sock.receive(), max(0.01, end - loop.time()))
|
|
data = json.loads(msg.data)
|
|
if data.get("type") == kind:
|
|
return data
|
|
|
|
|
|
async def drain(sock, timeout=0.3):
|
|
import asyncio
|
|
out = []
|
|
try:
|
|
while True:
|
|
msg = await asyncio.wait_for(sock.receive(), timeout)
|
|
out.append(json.loads(msg.data))
|
|
except asyncio.TimeoutError:
|
|
return out
|
|
|
|
|
|
@contextlib.asynccontextmanager
|
|
async def rig_server(config=None, workspace=True, setup=None):
|
|
data = tempfile.mkdtemp()
|
|
user_root = tempfile.mkdtemp()
|
|
os.makedirs(os.path.join(user_root, "workflows"))
|
|
store = Store(data)
|
|
# Tests connect from 127.0.0.1: nothing may be trusted, or checks are skipped.
|
|
store.update_config({"api_protection": {"enabled": True, "trusted_ips": []}, **(config or {})})
|
|
ps = FakePromptServer()
|
|
hub = Hub(store, ps)
|
|
if workspace:
|
|
hub.workspace = Workspace(store, os.path.join(user_root, "workflows"))
|
|
routes.setup(ps, store, hub)
|
|
ps.app.add_routes(comfy_routes(user_root))
|
|
if setup:
|
|
setup(ps.app)
|
|
server = TestServer(ps.app, host="127.0.0.1")
|
|
await server.start_server()
|
|
rig = Rig(server, store, hub, ps, user_root)
|
|
# unsafe=True: keep cookies for the 127.0.0.1 test host (browsers do).
|
|
async with ClientSession(cookie_jar=CookieJar(unsafe=True)) as http:
|
|
rig.http = http
|
|
try:
|
|
yield rig
|
|
finally:
|
|
await server.close()
|