Restart ComfyUI automatically after plugin updates

tests/test_o1key_restart_frontend.mjs

@@ -0,0 +1,83 @@
+import assert from "node:assert/strict";
+import fs from "node:fs";
+import vm from "node:vm";
+
+const source = fs.readFileSync(new URL("../web/js/o1keyRestart.js", import.meta.url), "utf8")
+    .replace(/^export /gm, "");
+const requests = [];
+const responses = [];
+const reloads = [];
+const attributes = new Map();
+const icon = { className: "pi pi-refresh side-bar-button-icon" };
+const label = { textContent: "重启" };
+const button = {
+    disabled: false,
+    setAttribute(name, value) { attributes.set(name, value); },
+    querySelector(selector) {
+        return selector === ".side-bar-button-icon" ? icon : label;
+    },
+};
+const document = {
+    querySelector(selector) {
+        assert.equal(selector, "#o1key-restart-button");
+        return button;
+    },
+};
+const window = {
+    location: {
+        href: "http://127.0.0.1:8188/?workflow=test",
+        replace(url) { reloads.push(url); },
+    },
+    setTimeout() { return 1; },
+    clearTimeout() {},
+};
+const reply = (data, status = 200) => ({
+    ok: status < 400,
+    status,
+    async json() { return data; },
+});
+async function fetch(url, options) {
+    requests.push({ url, options });
+    const next = responses.shift();
+    assert.ok(next, `Unexpected request: ${url}`);
+    return next;
+}
+const context = vm.createContext({ document, window, fetch, AbortController, URL, Date });
+vm.runInContext(source, context);
+
+responses.push(
+    reply({ ready: true, boot_id: "old" }),
+    reply({ error: "restart rejected" }, 500),
+);
+await assert.rejects(context.startComfyUIRestart(), /restart rejected/);
+assert.equal(reloads.length, 0);
+assert.equal(button.disabled, false);
+assert.equal(attributes.get("aria-busy"), "false");
+assert.equal(label.textContent, "重启");
+assert.match(icon.className, /pi-refresh/);
+requests.length = 0;
+
+responses.push(
+    reply({ ready: true, boot_id: "old" }),
+    reply({ success: true, boot_id: "old" }),
+    reply({ ready: true, boot_id: "new" }),
+    reply({}),
+);
+await context.startComfyUIRestart();
+assert.equal(requests.length, 4);
+assert.match(requests[0].url, /^\/o1key\/restart\/status\?/);
+assert.equal(requests[1].url, "/o1key/restart");
+assert.equal(requests[1].options.method, "POST");
+assert.equal(requests[1].options.body, "{}");
+assert.match(requests[2].url, /^\/o1key\/restart\/status\?/);
+assert.match(requests[3].url, /^\/api\/system_stats\?/);
+assert.equal(reloads.length, 1);
+const reloadedUrl = new URL(reloads[0]);
+assert.equal(reloadedUrl.searchParams.get("workflow"), "test");
+assert.ok(reloadedUrl.searchParams.has("o1key_restart"));
+assert.equal(button.disabled, true);
+assert.equal(attributes.get("aria-busy"), "true");
+assert.equal(label.textContent, "重启中");
+await assert.rejects(context.startComfyUIRestart(), /正在重启/);
+assert.equal(requests.length, 4);
+assert.equal(reloads.length, 1);

tests/test_o1key_update_button.mjs

@@ -96,7 +96,8 @@ let observer;
 const requests = [];
 const responses = [];
 let fetchFailure = false;
-let dismissNotice;
+let restartCalls = 0;
+let restartFailure = false;
 const reply = (body, status = 200) => ({ ok: status < 400, status, async json() { return body; } });
 const app = { registerExtension(value) { extension = value; } };
 const api = {
@@ -111,11 +112,11 @@ class MutationObserver {
     observe() {}
     disconnect() { this.disconnected = true; }
 }
-const window = {
-    setTimeout(callback) { dismissNotice = callback; return 1; },
-    clearTimeout() { dismissNotice = null; },
-};
-vm.runInNewContext(source, { app, api, document, window, MutationObserver });
+async function startComfyUIRestart() {
+    restartCalls++;
+    if (restartFailure) throw new Error("restart failed");
+}
+vm.runInNewContext(source, { app, api, document, startComfyUIRestart, MutationObserver });
 
 extension.setup();
 assert.ok(observer);
@@ -135,29 +136,15 @@ toolbar.append(token, restart);
 observer.callback();
 
 const update = document.querySelector("#o1key-update-button");
-const detect = document.querySelector("#o1key-detect-button");
 assert.equal(toolbar.children[0], token);
 assert.equal(toolbar.children[1], update);
-assert.equal(toolbar.children[2], detect);
-assert.equal(toolbar.children[3], restart);
+assert.equal(toolbar.children[2], restart);
 assert.equal(update.querySelector(".side-bar-button-label").textContent, "更新");
 assert.match(update.querySelector(".side-bar-button-icon").className, /pi-download/);
-assert.equal(detect.querySelector(".side-bar-button-label").textContent, "检测");
-assert.match(detect.querySelector(".side-bar-button-icon").className, /pi-check-circle/);
+assert.equal(document.querySelector("#o1key-detect-button"), null);
 assert.ok(observer.disconnected);
 extension.setup();
-assert.equal(toolbar.children.length, 4);
-
-detect.listeners.get("click")();
-const notice = document.querySelector("#o1key-detect-notice");
-assert.equal(notice.children[1].textContent, "检测按钮已加载");
-assert.equal(notice.hidden, false);
-assert.equal(requests.length, 0);
-dismissNotice();
-assert.equal(notice.hidden, true);
-detect.listeners.get("click")();
-assert.equal(notice.hidden, false);
-assert.equal(document.body.children.filter((child) => child.id === "o1key-detect-notice").length, 1);
+assert.equal(toolbar.children.length, 3);
 
 responses.push(reply({ update_available: false }));
 update.listeners.get("click")();
@@ -212,22 +199,34 @@ await new Promise(setImmediate);
 assert.equal(dialog.open, true);
 assert.equal(dialog.dataset.state, "available");
 
-responses.push(reply({ updated: true, version: "abc1234", requirements_changed: false }));
+responses.push(reply({ updated: true, version: "abc1234", requirements_changed: true }));
 await primary.listeners.get("click")();
 assert.equal(requests.at(-1).path, "/o1key/update");
 assert.equal(requests.at(-1).options.method, "POST");
 assert.equal(requests.at(-1).options.headers["X-O1Key-Update"], "1");
 assert.match(status.textContent, /更新完成/);
-assert.match(suggestion.textContent, /重启 ComfyUI/);
+assert.match(suggestion.textContent, /联系技术支持/);
+assert.doesNotMatch(suggestion.textContent, /requirements\.txt/);
 assert.equal(primary.textContent, "完成");
 assert.equal(dialog.dataset.state, "success");
+assert.equal(restartCalls, 0);
 
 responses.push(reply({ update_available: true }));
 await dialog.checkForUpdate();
-responses.push(reply({ updated: true, requirements_changed: true }));
+restartFailure = true;
+responses.push(reply({ updated: true, requirements_changed: false }));
 await primary.listeners.get("click")();
-assert.match(suggestion.textContent, /联系技术支持/);
-assert.doesNotMatch(suggestion.textContent, /requirements\.txt/);
+assert.equal(restartCalls, 1);
+assert.equal(dialog.dataset.state, "restart_failed");
+assert.match(status.textContent, /自动重启失败/);
+assert.match(suggestion.textContent, /左侧“重启”按钮/);
+assert.equal(primary.textContent, "关闭");
+await primary.listeners.get("click")();
+assert.equal(dialog.open, false);
+responses.push(reply({ update_available: true }));
+update.listeners.get("click")();
+await new Promise(setImmediate);
+assert.equal(dialog.dataset.state, "available");
 
 const requestsBeforeError = requests.length;
 responses.push(reply({ code: "local_changes", error: "插件目录有未提交的代码修改。", suggestion: "请先提交" }, 409));
@@ -254,3 +253,19 @@ assert.equal(primary.disabled, false);
 assert.equal(head.children.length, 1);
 header.children[2].listeners.get("click")();
 assert.equal(dialog.open, false);
+
+fetchFailure = false;
+restartFailure = false;
+responses.push(reply({ update_available: true }));
+update.listeners.get("click")();
+await new Promise(setImmediate);
+responses.push(reply({ updated: true, requirements_changed: false }));
+await primary.listeners.get("click")();
+assert.equal(restartCalls, 2);
+assert.equal(dialog.dataset.state, "restarting");
+assert.match(status.textContent, /正在重启 ComfyUI/);
+assert.match(suggestion.textContent, /页面会自动刷新/);
+assert.equal(primary.disabled, true);
+assert.equal(secondary.disabled, true);
+assert.equal(header.children[2].disabled, true);
+assert.equal(primary.textContent, "重启中…");