Restart ComfyUI automatically after plugin updates

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, "重启中…");