From 45a866f1949e98374ebbf6839b9e14a38e1595a7 Mon Sep 17 00:00:00 2001
From: Chenlei Hu <chenlei.hu@mail.utoronto.ca>
Date: Mon, 16 Sep 2024 17:14:11 +0900
Subject: [PATCH] Fix ComfyUI class setup procedure (#846)

---
 src/scripts/ui.ts            | 2 +-
 src/scripts/ui/menu/index.ts | 8 ++++++--
 tests-ui/utils/index.ts      | 1 +
 3 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/src/scripts/ui.ts b/src/scripts/ui.ts
index 8bb4e97ff..af40d92ae 100644
--- a/src/scripts/ui.ts
+++ b/src/scripts/ui.ts
@@ -373,7 +373,7 @@ export class ComfyUI {
     })
 
     // For testing. Legacy ui tests don't have vue app initialized.
-    if (!app.vueAppReady) {
+    if (window['IS_TEST']) {
       this.setup(document.body)
     }
   }
diff --git a/src/scripts/ui/menu/index.ts b/src/scripts/ui/menu/index.ts
index a9f3b7749..b753ccc26 100644
--- a/src/scripts/ui/menu/index.ts
+++ b/src/scripts/ui/menu/index.ts
@@ -181,7 +181,9 @@ export class ComfyAppMenu {
             resizeHandler = null
           }
           document.body.style.removeProperty('display')
-          app.ui.menuContainer.style.removeProperty('display')
+          if (app.ui.menuContainer) {
+            app.ui.menuContainer.style.removeProperty('display')
+          }
           this.element.style.display = 'none'
           app.ui.restoreMenuPosition()
         }
@@ -192,7 +194,9 @@ export class ComfyAppMenu {
 
   updatePosition(v: MenuPosition) {
     document.body.style.display = 'grid'
-    this.app.ui.menuContainer.style.display = 'none'
+    if (this.app.ui.menuContainer) {
+      this.app.ui.menuContainer.style.display = 'none'
+    }
     this.element.style.removeProperty('display')
     this.position = v
     if (v === 'Bottom') {
diff --git a/tests-ui/utils/index.ts b/tests-ui/utils/index.ts
index c4c38c407..f04ddedf1 100644
--- a/tests-ui/utils/index.ts
+++ b/tests-ui/utils/index.ts
@@ -38,6 +38,7 @@ export async function start(config: StartConfig = {}): Promise<StartResult> {
 
   Object.assign(localStorage, config.localStorage ?? {})
   document.body.innerHTML = html.toString()
+  window['IS_TEST'] = true
 
   mockApi(config)
   mockSettingStore()