Test suite: pytest unit and integration tests, node tests, CI workflow

This commit is contained in:
Doudou
2026-09-26 02:17:22 +00:00
parent 82b12c360d
commit 7a5385c46a
14 changed files with 1055 additions and 9 deletions
+1
View File
@@ -2,3 +2,4 @@
docs/
.github/
.gitattributes
tests/
+23
View File
@@ -0,0 +1,23 @@
name: Tests
on:
push:
pull_request:
jobs:
test:
name: Python and JavaScript tests
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v7
- uses: actions/setup-python@v6
with:
python-version: "3.11"
- uses: actions/setup-node@v5
with:
node-version: "22"
- name: Install test dependencies
run: pip install pytest aiohttp
- name: Python tests
run: python -m pytest -q
- name: JavaScript tests
run: node --test "tests/js/*.test.mjs"
+1
View File
@@ -96,6 +96,7 @@ Behind a reverse proxy, turn on WebSocket support, or nothing will be live. See
| [Settings](docs/settings.md) | Panel settings, ComfyUI settings, every key in `config.json` |
| [Reverse proxy](docs/reverse-proxy.md) | Nginx Proxy Manager, plain nginx, HTTPS |
| [Security](docs/security.md) | How passwords, sessions and keys are protected |
| [Development](docs/development.md) | Running the tests, code layout |
| [Changelog](CHANGELOG.md) | What changed, and when |
## Why this exists
+12 -9
View File
@@ -1,18 +1,16 @@
import logging
import os
import folder_paths
from server import PromptServer
from .rigshare_core.store import Store
from .rigshare_core.hub import Hub
from .rigshare_core.workspace import Workspace
from .rigshare_core import routes
NODE_CLASS_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS = {}
WEB_DIRECTORY = "./web"
try:
import folder_paths
from server import PromptServer
except ImportError: # imported outside ComfyUI, e.g. by the test runner
folder_paths = PromptServer = None
def data_folder():
return os.environ.get("RIGSHARE_DATA_DIR") or os.path.join(folder_paths.base_path, "rigshare")
@@ -27,8 +25,13 @@ def workflows_dir():
return os.path.join(user_root or os.path.join(folder_paths.get_user_directory(), "default"), "workflows")
if hasattr(PromptServer, "instance"):
if PromptServer is not None and hasattr(PromptServer, "instance"):
try:
from .rigshare_core.store import Store
from .rigshare_core.hub import Hub
from .rigshare_core.workspace import Workspace
from .rigshare_core import routes
folder = data_folder()
os.makedirs(folder, exist_ok=True)
store = Store(folder)
+37
View File
@@ -0,0 +1,37 @@
# Development
## Tests
RigShare has two test suites. Neither needs ComfyUI.
```bash
pip install pytest aiohttp
python -m pytest -q
node --test "tests/js/*.test.mjs"
```
| Suite | What it covers |
|---|---|
| `tests/python/test_store.py` | Accounts, first-run setup, password hashing, JWT sessions (tampering, `alg: none`, expiry, refresh, revocation), API keys, snapshots |
| `tests/python/test_graphdoc.py` | Workflow patches, link repair, id renumbering for simultaneous edits, and parity with the JavaScript implementation |
| `tests/python/test_workspace.py` | Folder rules: path normalisation (double encoding, `../`), private and shared folders, access lists, moves, deletes, browsing |
| `tests/python/test_server.py` | End to end against a real aiohttp server: login gate, cookies, API keys, permissions, Manager gating, template filtering, live rooms, access lists, folders through ComfyUI's userdata API |
| `tests/js/graphdoc.test.mjs` | The client-side patch logic |
| `tests/js/sync.test.mjs` | Client diffs, and the regression test for opening a file overwriting the previous tab's shared workflow |
The server tests use a small stand-in for ComfyUI's `/userdata` API (`tests/python/helpers.py`), so they run anywhere Python does. GitHub Actions runs both suites on every push and pull request.
`tests/` is excluded from the Comfy Registry package by `.comfyignore`.
## Layout
| Path | |
|---|---|
| `__init__.py` | ComfyUI entry point: wires the store, hub, workspace and routes |
| `rigshare_core/store.py` | Accounts, sessions, API keys, config, rooms and snapshots on disk |
| `rigshare_core/jwt.py` | HS256 tokens |
| `rigshare_core/hub.py` | Websocket hub: presence, chat, cursors, live rooms |
| `rigshare_core/graphdoc.py` | Workflow patches (mirrored by `web/rigshare/graphdoc.js`) |
| `rigshare_core/workspace.py` | Private and shared folders |
| `rigshare_core/routes.py` | HTTP API and the middleware that enforces login and permissions |
| `web/rigshare/` | The frontend extension |
+3
View File
@@ -16,3 +16,6 @@ Documentation = "https://github.com/xiaodoudou/ComfyUI-RigShare#documentation"
PublisherId = "xiaodoudou"
DisplayName = "ComfyUI RigShare"
Icon = "https://raw.githubusercontent.com/xiaodoudou/ComfyUI-RigShare/master/web/rigshare/assets/icon.png"
[tool.pytest.ini_options]
testpaths = ["tests/python"]
+40
View File
@@ -0,0 +1,40 @@
// Workflow patch logic on the client (mirror of rigshare_core/graphdoc.py).
import { test } from "node:test";
import assert from "node:assert/strict";
import { applyPatch, reconcile, linkId } from "../../web/rigshare/graphdoc.js";
const node = (id) => ({ id, type: "T", inputs: [{ name: "a", link: null }, { name: "b", link: null }], outputs: [{ name: "o", links: [] }] });
const doc = (links = []) => reconcile({ last_node_id: 3, last_link_id: links.length, nodes: [1, 2, 3].map(node), links, groups: [], extra: {} });
test("upserts and removes nodes without mutating the input", () => {
const base = doc();
const moved = { ...structuredClone(base.nodes[0]), pos: [5, 5] };
const out = applyPatch(base, { nodes: [moved, node(4)], removed: [2] });
assert.deepEqual(out.nodes.map((n) => n.id), [1, 3, 4]);
assert.equal(base.nodes.length, 3);
});
test("links merge by id instead of replacing each other", () => {
const out = applyPatch(doc([[1, 1, 0, 2, 0, "X"]]), { links: { upsert: [[2, 3, 0, 2, 1, "X"]] } });
assert.deepEqual(out.links.map(linkId).sort(), [1, 2]);
assert.deepEqual(out.nodes[1].inputs.map((s) => s.link), [1, 2]);
});
test("last ids never go backwards", () => {
const base = { ...doc(), last_node_id: 10 };
assert.equal(applyPatch(base, { set: { last_node_id: 4 } }).last_node_id, 10);
});
test("reconcile drops dangling links and keeps the newest link per input", () => {
const d = doc();
d.links = [[1, 1, 0, 2, 0, "X"], [2, 3, 0, 2, 0, "X"], [3, 9, 0, 2, 1, "X"]];
reconcile(d);
assert.deepEqual(d.links.map(linkId), [2]);
assert.equal(d.nodes[1].inputs[0].link, 2);
assert.deepEqual(d.nodes[2].outputs[0].links, [2]);
});
test("full replaces the document", () => {
const out = applyPatch(doc(), { full: { nodes: [], links: [] } });
assert.deepEqual(out, { nodes: [], links: [] });
});
+79
View File
@@ -0,0 +1,79 @@
// Client sync: diffs round-trip, and loading a graph never leaks into the old tab's room.
import { test } from "node:test";
import assert from "node:assert/strict";
globalThis.window ??= new EventTarget();
// RoomSync polls with setInterval; let those timers not keep the test process alive.
const nativeSetInterval = globalThis.setInterval;
globalThis.setInterval = (...args) => { const t = nativeSetInterval(...args); t.unref?.(); return t; };
globalThis.performance ??= { now: () => Date.now() };
const { RoomSync, snapshotOf, diffSnapshots } = await import("../../web/rigshare/sync.js");
const { applyPatch, reconcile } = await import("../../web/rigshare/graphdoc.js");
const sleep = (ms) => new Promise((r) => setTimeout(r, ms));
const node = (id) => ({ id, type: "T", pos: [0, 0], inputs: [{ name: "a", link: null }, { name: "b", link: null }], outputs: [{ name: "o", links: [] }] });
test("diff then apply reproduces the edited workflow", () => {
let seed = 3;
const rnd = () => (seed = (seed * 16807) % 2147483647) / 2147483647;
for (let t = 0; t < 300; t++) {
const base = reconcile({ last_node_id: 5, last_link_id: 2, nodes: [1, 2, 3, 4, 5].map(node),
links: [[1, 1, 0, 2, 0, "X"], [2, 3, 0, 4, 1, "X"]], groups: [{ id: 1, title: "g" }], extra: { a: 1 }, version: 0.4 });
const cur = structuredClone(base);
if (rnd() < 0.5) cur.nodes[0].pos = [rnd(), rnd()];
if (rnd() < 0.4) { cur.last_node_id++; cur.nodes.push(node(cur.last_node_id)); }
if (rnd() < 0.4) { cur.last_link_id++; cur.links.push([cur.last_link_id, cur.nodes.at(-1).id, 0, 5, 0, "X"]); }
if (rnd() < 0.3) cur.links = cur.links.filter((l) => l[0] !== 1);
if (rnd() < 0.3) cur.groups[0].title = "renamed";
if (rnd() < 0.2) cur.extra.b = 2;
if (rnd() < 0.15) cur.nodes = cur.nodes.filter((n) => n.id !== 3);
reconcile(cur);
const patch = diffSnapshots(snapshotOf(base), snapshotOf(cur));
const out = patch ? applyPatch(base, patch) : base;
assert.equal(diffSnapshots(snapshotOf(out), snapshotOf(cur)), null);
}
});
test("opening a file never overwrites the shared workflow of the previous tab", async () => {
const graphA = { nodes: [node(1), node(2)], links: [], extra: {}, version: 0.4 };
const fromDisk = { nodes: [{ id: 1, type: "ReActorFaceSwap" }, { id: 2, type: "LoadImage" }, { id: 3, type: "SaveImage" }], links: [], extra: {}, version: 0.4 };
let current = structuredClone(graphA);
const tabA = { path: "workflows/shared.json", filename: "shared.json", isTemporary: false };
const store = { activeWorkflow: tabA, openWorkflows: [tabA] };
const sent = [];
const client = Object.assign(new EventTarget(), {
online: true, perms: { edit: true }, rooms: [], users: [],
on(type, fn) { this.addEventListener(type, (e) => fn(e.detail)); },
emit(type, detail) { this.dispatchEvent(new CustomEvent(type, { detail })); },
send(m) { sent.push(m); return true; },
});
const app = {
extensionManager: { workflow: store }, canvas: null,
rootGraph: { serialize: () => structuredClone(current), extra: {}, getNodeById: () => null },
// Like ComfyUI: swap the graph, yield, then switch to a new tab.
async loadGraphData(data, clean, restore, target) {
current = structuredClone(data);
await sleep(30);
if (typeof target === "string" || target == null) {
const tab = { path: `workflows/${target}`, filename: target, isTemporary: true };
store.openWorkflows.push(tab);
store.activeWorkflow = tab;
}
},
};
const api = new EventTarget();
const sync = new RoomSync(app, api, client);
sync.overrides[tabA.path] = true;
sync.tick();
client.emit("room", { room: `file:${tabA.path}`, version: 1, doc: graphA, role: "edit", manage: true });
await sleep(400);
sent.length = 0;
const loading = app.loadGraphData(fromDisk, true, true, "face swap.json");
for (let i = 0; i < 5; i++) { api.dispatchEvent(new Event("graphChanged")); sync.checkLocal(); sync.tick(); await sleep(5); }
await loading;
for (let i = 0; i < 5; i++) { api.dispatchEvent(new Event("graphChanged")); sync.tick(); await sleep(20); }
const leaked = sent.filter((m) => m.type === "graph" && m.room === `file:${tabA.path}`);
assert.equal(leaked.length, 0, "the previous tab's room must not receive the new graph");
});
+21
View File
@@ -0,0 +1,21 @@
import asyncio
import os
import sys
import pytest
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
@pytest.fixture
def run():
"""Run a coroutine to completion (no pytest-asyncio dependency)."""
def _run(coro):
return asyncio.run(coro)
return _run
@pytest.fixture
def store(tmp_path):
from helpers import Store
return Store(str(tmp_path))
+199
View File
@@ -0,0 +1,199 @@
"""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()
+150
View File
@@ -0,0 +1,150 @@
"""Workflow patches: apply, link repair, conflict renumbering, JS parity."""
import copy
import json
import random
import shutil
import subprocess
from pathlib import Path
import pytest
from helpers import ROOT
from rigshare_core.graphdoc import apply_patch, reconcile, remap_conflicts, link_id, link_ends
def node(nid, n_in=2, n_out=1):
return {"id": nid, "type": "T", "pos": [0, 0],
"inputs": [{"name": f"i{i}", "link": None} for i in range(n_in)],
"outputs": [{"name": f"o{j}", "links": []} for j in range(n_out)]}
def doc_with(n=3, links=()):
doc = {"last_node_id": n, "last_link_id": len(links), "nodes": [node(i + 1) for i in range(n)],
"links": [list(l) for l in links], "groups": [], "extra": {}, "version": 0.4}
return reconcile(doc)
def assert_consistent(doc):
ids = {n["id"] for n in doc["nodes"]}
links = {link_id(l): l for l in doc["links"]}
assert len(links) == len(doc["links"]), "duplicate link ids"
assert len(ids) == len(doc["nodes"]), "duplicate node ids"
for l in doc["links"]:
origin, _, target, _ = link_ends(l)
assert origin in ids and target in ids, "dangling link"
for n in doc["nodes"]:
for i, slot in enumerate(n["inputs"]):
if slot["link"] is not None:
_, _, target, target_slot = link_ends(links[slot["link"]])
assert (target, target_slot) == (n["id"], i)
def test_upsert_and_remove_nodes():
doc = doc_with(3)
moved = copy.deepcopy(doc["nodes"][0]); moved["pos"] = [5, 5]
out = apply_patch(doc, {"nodes": [moved, node(4)], "removed": [2]})
assert [n["id"] for n in out["nodes"]] == [1, 3, 4]
assert out["nodes"][0]["pos"] == [5, 5]
assert doc["nodes"][0]["pos"] == [0, 0], "input document must not be mutated"
def test_links_merge_by_id():
doc = doc_with(3, [(1, 1, 0, 2, 0, "X")])
out = apply_patch(doc, {"links": {"upsert": [[2, 3, 0, 2, 1, "X"]]}})
assert sorted(link_id(l) for l in out["links"]) == [1, 2], "a new link must not erase another"
def test_max_keys_never_go_backwards():
doc = doc_with(3)
doc["last_node_id"] = 10
assert apply_patch(doc, {"set": {"last_node_id": 4}})["last_node_id"] == 10
def test_reconcile_drops_dangling_and_duplicate_inputs():
doc = doc_with(3)
doc["links"] = [[1, 1, 0, 2, 0, "X"], [2, 3, 0, 2, 0, "X"], [3, 9, 0, 2, 1, "X"]]
reconcile(doc)
assert [link_id(l) for l in doc["links"]] == [2], "newest link on a slot wins, dangling dropped"
assert doc["nodes"][1]["inputs"][0]["link"] == 2
assert doc["nodes"][2]["outputs"][0]["links"] == [2]
assert_consistent(doc)
def test_remap_simultaneous_adds():
base = doc_with(2)
def add(target_slot):
return {"nodes": [node(3)], "added": [3],
"links": {"upsert": [[1, 3, 0, 2, target_slot, "X"]], "added": [1]},
"set": {"last_node_id": 3, "last_link_id": 1}}
server = base
pa, changed_a = remap_conflicts(server, add(0)); server = apply_patch(server, pa)
pb, changed_b = remap_conflicts(server, add(1)); server = apply_patch(server, pb)
assert not changed_a and changed_b
assert sorted(n["id"] for n in server["nodes"]) == [1, 2, 3, 4]
assert sorted(link_id(l) for l in server["links"]) == [1, 2]
assert_consistent(server)
def test_remap_leaves_updates_of_existing_nodes_alone():
doc = doc_with(3)
update = copy.deepcopy(doc["nodes"][0]); update["pos"] = [1, 1]
patch, changed = remap_conflicts(doc, {"nodes": [update]})
assert not changed and patch["nodes"][0]["id"] == 1
def test_concurrent_edits_converge():
rnd = random.Random(1)
for _ in range(100):
server = doc_with(4, [(1, 1, 0, 2, 0, "X")])
observer = copy.deepcopy(server)
patches = []
for _ in range(2): # two people, same base, same instant
nid, lid = server["last_node_id"] + 1, server["last_link_id"] + 1
target = rnd.choice([n["id"] for n in server["nodes"]])
patches.append({"nodes": [node(nid)], "added": [nid],
"links": {"upsert": [[lid, nid, 0, target, rnd.randint(0, 1), "X"]], "added": [lid]},
"set": {"last_node_id": nid, "last_link_id": lid}})
for p in patches:
p2, _ = remap_conflicts(server, p)
server = apply_patch(server, p2)
observer = apply_patch(observer, p2)
assert_consistent(server)
assert len(server["nodes"]) == 6
assert json.dumps(server, sort_keys=True) == json.dumps(observer, sort_keys=True)
@pytest.mark.skipif(not shutil.which("node"), reason="node is not installed")
def test_javascript_mirror_matches(tmp_path):
"""graphdoc.js must produce exactly the same documents as graphdoc.py."""
rnd = random.Random(7)
cases = []
for _ in range(60):
doc = doc_with(5, [(1, 1, 0, 2, 0, "X"), (2, 3, 0, 4, 1, "X")])
patches, cur = [], doc
for _ in range(5):
ids = [n["id"] for n in cur["nodes"]]
p = {}
if rnd.random() < .5:
n = copy.deepcopy(rnd.choice(cur["nodes"])); n["pos"] = [rnd.randint(0, 99), 0]; p["nodes"] = [n]
if rnd.random() < .4 and len(ids) > 1:
a, b = rnd.sample(ids, 2); lid = cur["last_link_id"] + 1
p["links"] = {"upsert": [[lid, a, 0, b, rnd.randint(0, 1), "X"]], "added": [lid]}
p["set"] = {"last_link_id": lid}
if rnd.random() < .2 and len(ids) > 3:
p["removed"] = [rnd.choice(ids)]
patches.append(p)
cur = apply_patch(cur, p)
cases.append({"doc": doc, "patches": patches, "expected": cur})
data = tmp_path / "cases.json"
data.write_text(json.dumps(cases))
module = (Path(ROOT) / "web" / "rigshare" / "graphdoc.js").as_uri()
script = tmp_path / "check.mjs"
script.write_text(
f'import fs from "fs"; import {{ applyPatch }} from "{module}";\n'
'const cases = JSON.parse(fs.readFileSync(process.argv[2], "utf8")); let bad = 0;\n'
'for (const c of cases) { let cur = c.doc; for (const p of c.patches) cur = applyPatch(cur, p);\n'
' if (JSON.stringify(cur) !== JSON.stringify(c.expected)) bad++; }\n'
'console.log(bad); process.exit(bad ? 1 : 0);\n')
result = subprocess.run(["node", str(script), str(data)], capture_output=True, text=True)
assert result.returncode == 0, f"{result.stdout} mismatching cases\n{result.stderr}"
+199
View File
@@ -0,0 +1,199 @@
"""End-to-end: HTTP API, login gate, websocket hub, folders over ComfyUI's API."""
import asyncio
import json
import os
from urllib.parse import quote
from helpers import PASSWORD, enc, rig_server
BASE_DOC = {"last_node_id": 2, "last_link_id": 0, "nodes": [
{"id": 1, "type": "A", "inputs": [], "outputs": [{"name": "o", "links": []}]},
{"id": 2, "type": "B", "inputs": [{"name": "i", "link": None}, {"name": "j", "link": None}], "outputs": []}],
"links": [], "groups": [], "extra": {}, "version": 0.4}
def test_login_gate_and_setup(run):
async def scenario():
async with rig_server() as rig:
h = rig.http
r = await h.get(rig.url("/"), headers={"Accept": "text/html"}, allow_redirects=False)
assert r.status == 302 and "/rigshare/login" in r.headers["Location"]
assert (await h.get(rig.url("/api/prompt"))).status == 401
status = await (await h.get(rig.url("/rigshare/api/status"))).json()
assert status["needs_setup"]
page = await (await h.get(rig.url("/rigshare/login"))).text()
assert "const SETUP = true" in page
await rig.setup_admin()
r = await h.post(rig.url("/rigshare/api/setup"), json={"username": "evil", "password": PASSWORD})
assert r.status == 409, "setup works only once"
r = await h.post(rig.url("/rigshare/api/login"), json={"username": "boss", "password": "nope-nope"})
assert r.status == 401
run(scenario())
def test_cookie_session_and_logout(run):
async def scenario():
async with rig_server() as rig:
await rig.setup_admin()
r = await rig.http.post(rig.url("/rigshare/api/login"), json={"username": "boss", "password": PASSWORD})
token = (await r.json())["token"]
assert "rigshare_session" in r.cookies
assert (await rig.http.get(rig.url("/api/prompt"))).status == 200, "cookie authenticates"
await rig.http.post(rig.url("/rigshare/api/logout"), headers={"Authorization": f"Bearer {token}"})
assert (await rig.http.get(rig.url("/api/prompt"), headers={"Authorization": f"Bearer {token}"})).status == 401
run(scenario())
def test_api_keys_over_http(run):
async def scenario():
async with rig_server() as rig:
admin = await rig.setup_admin()
r = await rig.http.post(rig.url("/rigshare/api/me/keys"), json={"name": "script"}, headers=admin)
body = await r.json()
key = {"Authorization": f"Bearer {body['key']}"}
assert (await rig.http.get(rig.url("/api/prompt"), headers=key)).status == 200
await rig.http.delete(rig.url(f"/rigshare/api/me/keys/{body['id']}"), headers=admin)
assert (await rig.http.get(rig.url("/api/prompt"), headers=key)).status == 401
run(scenario())
def test_permissions(run):
async def scenario():
async with rig_server() as rig:
admin = await rig.setup_admin()
viewer = await rig.add_user("viewer")
queuer = await rig.add_user("queuer", queue=True)
h = rig.http
assert (await h.post(rig.url("/api/prompt"), headers=viewer)).status == 403
assert (await h.post(rig.url("/api/prompt"), headers=queuer)).status == 200
assert (await h.post(rig.url("/api/manager/queue/install"), headers=queuer)).status == 403
assert (await h.get(rig.url("/customnode/getmappings"), headers=viewer)).status == 200, "read-only lookup"
assert (await h.post(rig.url("/api/manager/queue/install"), headers=admin)).status == 200
await h.patch(rig.url("/rigshare/api/users/queuer"), json={"perms": {"manager": True}}, headers=admin)
assert (await h.post(rig.url("/api/manager/queue/install"), headers=queuer)).status == 200
assert (await h.get(rig.url("/rigshare/api/users"), headers=viewer)).status == 403
run(scenario())
def test_template_filter_and_live_options(run, tmp_path):
templates = tmp_path / "templates"
templates.mkdir()
(templates / "index.json").write_text(json.dumps([
{"title": "Image", "templates": [{"name": "sdxl", "tags": ["Image"], "openSource": True},
{"name": "api_flux", "tags": ["API"], "openSource": False}]},
{"title": "Ads", "templates": [{"name": "x", "tags": ["API"]}]}]))
async def scenario():
async with rig_server(setup=lambda app: app.router.add_static("/templates", str(templates))) as rig:
admin = await rig.setup_admin()
alice = await rig.add_user("alice", edit=True)
sock = await rig.ws("alice")
full = await (await rig.http.get(rig.url("/templates/index.json"), headers=alice)).json()
assert sum(len(c["templates"]) for c in full) == 3
await rig.http.patch(rig.url("/rigshare/api/config"), json={"hide_api_templates": True, "hide_comfy_account": True}, headers=admin)
assert (await sock.until("server"))["server"]["hide_comfy_account"] is True
slim = await (await rig.http.get(rig.url("/templates/index.json"), headers=alice)).json()
assert [t["name"] for c in slim for t in c["templates"]] == ["sdxl"] and len(slim) == 1
await sock.close()
run(scenario())
def test_rooms_conflicts_and_access_lists(run):
async def scenario():
async with rig_server() as rig:
admin = await rig.setup_admin()
for name in ("alice", "bob", "carol"):
await rig.add_user(name, edit=True)
alice, bob, carol = [await rig.ws(n) for n in ("alice", "bob", "carol")]
room = "file:workflows/w.json"
await alice.send({"type": "join", "room": room, "name": "w", "doc": BASE_DOC})
joined = await alice.until("room")
assert joined["created"] and joined["manage"]
await bob.send({"type": "join", "room": room, "name": "w", "doc": BASE_DOC})
await bob.until("room")
await alice.drain(); await bob.drain()
def add(target_slot):
return {"nodes": [{"id": 3, "type": "N", "inputs": [], "outputs": [{"name": "o", "links": [1]}]}], "added": [3],
"links": {"upsert": [[1, 3, 0, 2, target_slot, "X"]], "added": [1]},
"set": {"last_node_id": 3, "last_link_id": 1}}
await alice.send({"type": "graph", "room": room, "patch": add(0)})
await bob.send({"type": "graph", "room": room, "patch": add(1)})
await asyncio.sleep(0.4)
doc = rig.hub.rooms[room].doc
assert sorted(n["id"] for n in doc["nodes"]) == [1, 2, 3, 4], "both nodes kept"
assert [s["link"] for s in doc["nodes"][1]["inputs"]] == [1, 2]
assert any(m["type"] == "doc" and m.get("reason") is None for m in await bob.drain()), "quiet resync"
r = await rig.http.put(rig.url("/rigshare/api/room/acl"), json={"key": room, "restricted": True, "members": {"carol": "view"}}, headers=rig.headers["bob"])
assert r.status == 403, "only owner or admin"
r = await rig.http.put(rig.url("/rigshare/api/room/acl"), json={"key": room, "restricted": True, "members": {"carol": "view"}}, headers=rig.headers["alice"])
assert r.status == 200
msgs = await bob.drain()
assert any(m["type"] == "room_denied" for m in msgs)
presence = [m for m in msgs if m["type"] == "presence"][-1]
assert all(r["key"] != room for r in presence["rooms"])
assert (await rig.http.get(rig.url(f"/api/userdata/{enc('workflows/w.json')}"), headers=rig.headers["bob"])).status == 403
await carol.send({"type": "join", "room": room, "name": "w", "doc": BASE_DOC})
assert (await carol.until("room"))["role"] == "view"
await carol.drain()
await carol.send({"type": "graph", "room": room, "patch": {"nodes": [dict(BASE_DOC["nodes"][0], pos=[9, 9])]}})
assert "view" in ((await carol.until("doc")).get("reason") or "")
r = await rig.http.put(rig.url("/rigshare/api/room/acl"), json={"key": room, "restricted": True, "members": {"carol": "edit"}}, headers=admin)
assert (await carol.until("room_info"))["role"] == "edit"
for s in (alice, bob, carol):
await s.close()
run(scenario())
def test_folders_over_comfy_api(run):
async def scenario():
async with rig_server() as rig:
await rig.setup_admin()
for name, perms in (("alice", {"edit": True}), ("bob", {"edit": True}), ("carol", {"edit": True}), ("dave", {})):
await rig.add_user(name, **perms)
H, h = rig.headers, rig.http
async def save(user, path):
return (await h.post(rig.url(f"/api/userdata/{enc(path)}"), data=b"{}", headers=H[user])).status
async def listing(user):
r = await h.get(rig.url("/api/userdata?dir=workflows&recurse=true&full_info=true"), headers=H[user])
return {i["path"] for i in await r.json()}
assert await save("alice", "workflows/users/alice/secret.json") == 200
assert await save("dave", "workflows/users/dave/mine.json") == 200
assert await save("dave", "workflows/common.json") == 403
assert await save("bob", "workflows/users/alice/evil.json") == 403
for raw in (enc("workflows/users/alice/secret.json"),
quote(enc("workflows/users/alice/secret.json"), safe=""),
enc("workflows/shared/../users/alice/secret.json"),
"workflows/users/alice/secret.json"):
assert (await h.get(rig.url(f"/api/userdata/{raw}"), headers=H["bob"])).status in (403, 404), raw
assert "users/alice/secret.json" not in await listing("bob")
assert "users/alice/secret.json" in await listing("alice")
r = await h.post(rig.url("/rigshare/api/workspace/mkdir"), json={"parent": "@shared", "name": "Team"}, headers=H["bob"])
assert r.status == 200
assert await save("alice", "workflows/shared/Team/plan.json") == 200
await h.patch(rig.url("/rigshare/api/workspace/folders"), json={"path": "shared/Team", "restricted": True, "members": {"carol": "view"}}, headers=H["bob"])
assert not any(p.startswith("shared/Team/") for p in await listing("alice"))
assert await save("carol", "workflows/shared/Team/plan.json") == 403
assert (await h.get(rig.url(f"/api/userdata/{enc('workflows/shared/Team/plan.json')}"), headers=H["carol"])).status == 200
r = await h.post(rig.url(f"/api/userdata/{enc('workflows/shared/Team/plan.json')}/move/{enc('workflows/users/bob/plan.json')}"), headers=H["bob"])
assert r.status == 200
assert (await h.delete(rig.url(f"/api/userdata/{enc('workflows/users/alice/secret.json')}"), headers=H["bob"])).status == 403
assert (await h.delete(rig.url(f"/api/userdata/{enc('workflows/users/alice/secret.json')}"), headers=H["alice"])).status == 204
sock = await rig.ws("bob")
await sock.send({"type": "join", "room": "file:workflows/users/alice/other.json", "name": "o", "doc": {"nodes": []}})
assert (await sock.until("room_denied"))["room"].endswith("alice/other.json")
await sock.close()
await save("alice", "workflows/users/alice/keep.json")
await h.patch(rig.url("/rigshare/api/users/alice"), json={"username": "alicia"}, headers=H["boss"])
assert os.path.exists(os.path.join(rig.user_root, "workflows", "users", "alicia", "keep.json"))
run(scenario())
+157
View File
@@ -0,0 +1,157 @@
"""Accounts, passwords, JWT sessions, API keys, snapshots."""
import json
import time
import pytest
from helpers import PASSWORD
from rigshare_core import jwt
import rigshare_core.store as store_module
def make_admin(store, name="admin"):
return store.setup_admin(name, PASSWORD)
# ----- setup and accounts ------------------------------------------------
def test_no_default_accounts(store):
assert store.needs_setup
assert store.users == {}
def test_setup_only_once(store):
user = make_admin(store)
assert user["perms"]["admin"] and not store.needs_setup
with pytest.raises(PermissionError):
store.setup_admin("intruder", PASSWORD)
def test_env_provisioning(tmp_path, monkeypatch):
monkeypatch.setenv("RIGSHARE_ADMIN_USER", "ops")
monkeypatch.setenv("RIGSHARE_ADMIN_PASSWORD", PASSWORD)
s = store_module.Store(str(tmp_path))
assert s.users["ops"]["perms"]["admin"]
def test_passwords_are_hashed(store, tmp_path):
make_admin(store)
raw = (tmp_path / "users.json").read_text()
assert PASSWORD not in raw
assert store.users["admin"]["password"].startswith("pbkdf2_sha256$")
def test_authenticate(store):
make_admin(store)
assert store.authenticate("admin", PASSWORD)
assert store.authenticate("admin", "wrong-password") is None
assert store.authenticate("nobody", PASSWORD) is None
@pytest.mark.parametrize("username,password", [("x", PASSWORD), ("bad name", PASSWORD), ("ok_name", "short")])
def test_create_user_validation(store, username, password):
make_admin(store)
with pytest.raises(ValueError):
store.create_user(username, password)
def test_last_admin_is_protected(store):
make_admin(store)
with pytest.raises(ValueError):
store.delete_user("admin")
with pytest.raises(ValueError):
store.update_user("admin", perms={"admin": False})
def test_password_change_revokes_sessions(store):
make_admin(store)
token = store.create_session("admin")
store.update_user("admin", password="another-password")
assert store.user_for_token(token) is None
# ----- JWT sessions --------------------------------------------------------
def claims_of(token):
return json.loads(jwt._b64d(token.split(".")[1]))
def test_session_is_jwt(store):
make_admin(store)
token = store.create_session("admin")
assert token.count(".") == 2
assert store.user_for_token(token)["username"] == "admin"
def test_tampered_token_rejected(store):
make_admin(store)
header, body, sig = store.create_session("admin").split(".")
forged = jwt._b64e(json.dumps({**claims_of(f"{header}.{body}.{sig}"), "sub": "someone"}).encode())
assert store.user_for_token(f"{header}.{forged}.{sig}") is None
assert store.user_for_token(f"{header}.{body}.{sig[:-2]}AA") is None
def test_alg_none_rejected(store):
make_admin(store)
_, body, _ = store.create_session("admin").split(".")
none_header = jwt._b64e(b'{"alg":"none","typ":"JWT"}')
assert store.user_for_token(f"{none_header}.{body}.") is None
def test_expired_token_rejected(store):
make_admin(store)
sid = claims_of(store.create_session("admin"))["sid"]
old = jwt.encode({"iss": "rigshare", "sub": "admin", "sid": sid, "iat": 0, "nbf": 0, "exp": 1}, store.secret)
assert store.user_for_token(old) is None
def test_logout_revokes_immediately(store):
make_admin(store)
token = store.create_session("admin")
store.revoke_session(token)
assert store.user_for_token(token) is None
def test_refresh_only_when_old(store, monkeypatch):
make_admin(store)
token = store.create_session("admin")
assert store.refresh_token(token) is None
monkeypatch.setattr(store_module, "TOKEN_REFRESH", -1)
fresh = store.refresh_token(token)
assert fresh and store.user_for_token(fresh)["username"] == "admin"
assert claims_of(fresh)["sid"] == claims_of(token)["sid"]
def test_rename_keeps_sessions_and_keys(store):
make_admin(store)
token = store.create_session("admin")
key, _ = store.create_api_key("admin", "k")
store.rename_user("admin", "root")
assert store.user_for_token(token)["username"] == "root"
assert store.user_for_token(key)["username"] == "root"
# ----- API keys ---------------------------------------------------------------
def test_api_keys(store, tmp_path):
make_admin(store)
key, meta = store.create_api_key("admin", "script")
assert key.startswith("rs_") and key not in (tmp_path / "users.json").read_text()
assert store.user_for_token(key)["username"] == "admin"
assert "hash" not in store.list_api_keys("admin")[0]
store.revoke_api_key("admin", meta["id"])
assert store.user_for_token(key) is None
# ----- snapshots ----------------------------------------------------------------
def test_snapshot_rotation_keeps_labelled(store):
store.update_config({"snapshot_keep": 2})
for i in range(4):
store.add_snapshot({"nodes": []}, "me", "room", "Room")
time.sleep(0.01)
store.add_snapshot({"nodes": []}, "me", "room", "Room", label="keep me")
snaps = store.list_snapshots("room")
assert len([s for s in snaps if not s["label"]]) == 2
assert any(s["label"] == "keep me" for s in snaps)
+133
View File
@@ -0,0 +1,133 @@
"""Folder rules: path normalisation, roles, deletes, moves, browsing."""
import os
import pytest
from helpers import Store
from rigshare_core.workspace import Workspace, norm
ADMIN = {"edit": True, "queue": True, "admin": True}
EDITOR = {"edit": True, "queue": True, "admin": False}
VIEWER = {"edit": False, "queue": False, "admin": False}
@pytest.fixture
def ws(tmp_path):
root = tmp_path / "workflows"
root.mkdir()
return Workspace(Store(str(tmp_path / "data")), str(root))
def touch(ws, rel):
path = os.path.join(ws.root, *rel.split("/"))
os.makedirs(os.path.dirname(path), exist_ok=True)
open(path, "w").write("{}")
@pytest.mark.parametrize("raw,expected", [
("workflows/a.json", "workflows/a.json"),
("workflows%2Fusers%2Fbob%2Fa.json", "workflows/users/bob/a.json"),
("workflows%252Fusers%252Fbob%252Fa.json", "workflows/users/bob/a.json"),
("workflows/shared/../users/bob/a.json", "workflows/users/bob/a.json"),
("../../etc/passwd", "etc/passwd"),
("workflows\\users\\bob", "workflows/users/bob"),
("", ""),
])
def test_norm(raw, expected):
assert norm(raw) == expected
def test_areas():
assert Workspace.area("users/bob/x.json") == ("private", "bob")
assert Workspace.area("shared/Team/x.json") == ("shared", "shared/Team")
assert Workspace.area("x.json") == ("common", None)
assert Workspace.workflows_rel("workflows/a/b.json") == "a/b.json"
assert Workspace.workflows_rel("subgraphs/x.json") is None
def test_private_space(ws):
assert ws.role("users/bob/x.json", "user", "bob", VIEWER) == "edit", "own space even for viewer accounts"
assert ws.role("users/bob/x.json", "user", "alice", EDITOR) is None
assert ws.role("users/bob/x.json", "user", "root", ADMIN) == "edit"
assert ws.can_delete("users/bob/x.json", "user", "bob", VIEWER)
assert not ws.can_delete("users/bob/x.json", "user", "alice", EDITOR)
def test_common_space(ws):
assert ws.role("x.json", "user", "alice", EDITOR) == "edit"
assert ws.role("x.json", "user", "dave", VIEWER) == "view"
assert not ws.can_delete("x.json", "user", "alice", EDITOR), "Common deletes are admin-only"
assert ws.can_delete("x.json", "user", "root", ADMIN)
def test_shared_folder_access_list(ws):
folder = ws.create_folder("Team", "bob")
assert ws.role(f"{folder}/a.json", "user", "alice", EDITOR) == "edit", "open by default"
ws.set_acl(folder, {"members": {"carol": "view", "erin": "edit"}})
assert ws.role(f"{folder}/a.json", "user", "alice", EDITOR) is None
assert ws.role(f"{folder}/a.json", "user", "carol", EDITOR) == "view"
assert ws.role(f"{folder}/a.json", "user", "erin", EDITOR) == "edit"
assert ws.role(f"{folder}/a.json", "user", "erin", VIEWER) == "view", "a role never exceeds account perms"
assert ws.role(f"{folder}/a.json", "user", "bob", EDITOR) == "edit", "owner always"
assert ws.can_delete(f"{folder}/a.json", "user", "bob", EDITOR)
assert not ws.can_delete(f"{folder}/a.json", "user", "erin", EDITOR)
def test_moves(ws):
folder = ws.create_folder("Team", "bob")
assert ws.can_move("users/bob/a.json", f"{folder}/a.json", "user", "bob", EDITOR)
assert ws.can_move(f"{folder}/a.json", f"{folder}/b.json", "user", "alice", EDITOR), "rename in place"
assert not ws.can_move(f"{folder}/a.json", "users/alice/a.json", "user", "alice", EDITOR), "only owner moves out"
assert not ws.can_move("users/bob/a.json", "users/alice/a.json", "user", "bob", EDITOR)
@pytest.mark.parametrize("name", ["", "..", "a/b", "x" * 80, ".hidden"])
def test_bad_folder_names(ws, name):
with pytest.raises(ValueError):
ws.create_folder(name, "bob")
def test_list_dir(ws):
folder = ws.create_folder("Team", "bob")
touch(ws, "common.json")
touch(ws, "users/alice/secret.json")
touch(ws, f"{folder}/plan.json")
shared = ws.list_dir("@shared", "user", "alice", EDITOR)
assert {e["name"] for e in shared["entries"]} == {"Team", "common"}
ws.set_acl(folder, {"members": {}})
shared = ws.list_dir("@shared", "user", "alice", EDITOR)
assert {e["name"] for e in shared["entries"]} == {"common"}
with pytest.raises(PermissionError):
ws.list_dir("users/alice", "user", "bob", EDITOR)
with pytest.raises(PermissionError):
ws.list_dir("users", "user", "bob", EDITOR)
mine = ws.list_dir("users/alice", "user", "alice", EDITOR)
assert [e["name"] for e in mine["entries"]] == ["secret"] and mine["can_mkdir"]
def test_folder_operations(ws):
touch(ws, "users/alice/keep.json")
sub = ws.make_dir("users/alice", "renders", "user", "alice", EDITOR)
assert sub == "users/alice/renders"
with pytest.raises(PermissionError):
ws.make_dir("users/alice", "x", "user", "bob", EDITOR)
with pytest.raises(PermissionError):
ws.rename_dir("users/alice", "hack", "user", "alice", EDITOR)
assert ws.rename_dir(sub, "finals", "user", "alice", EDITOR) == "users/alice/finals"
team = ws.make_dir("@shared", "Team", "user", "bob", EDITOR)
assert team == "shared/Team" and ws.folders[team]["owner"] == "bob"
assert ws.move_dir("users/alice/finals", team, "user", "alice", EDITOR) == "shared/Team/finals"
with pytest.raises(PermissionError):
ws.delete_dir(team, "user", "alice", EDITOR)
ws.delete_dir("shared/Team/finals", "user", "bob", EDITOR)
with pytest.raises(PermissionError):
ws.delete_dir("users/alice", "user", "alice", EDITOR)
def test_rename_user_moves_private_folder(ws):
touch(ws, "users/alice/keep.json")
ws.create_folder("Team", "alice")
ws.rename_user("alice", "alicia")
assert os.path.exists(os.path.join(ws.root, "users", "alicia", "keep.json"))
assert ws.folders["shared/Team"]["owner"] == "alicia"