Mock settingStore

Comfy-Org · Sep 16, 2024 · cb5e32d · cb5e32d
1 parent 74fa4a2
commit cb5e32d
Show file tree

Hide file tree

Showing 4 changed files with 38 additions and 27 deletions.
diff --git a/src/extensions/core/nodeBadge.ts b/src/extensions/core/nodeBadge.ts
@@ -84,10 +84,6 @@ class NodeBadgeExtension implements ComfyExtension {
   ) {}
 
   init(app: ComfyApp) {
-    if (!app.vueAppReady) {
-      return
-    }
-
     const settingStore = useSettingStore()
     this.nodeSourceBadgeMode = computed(
       () =>
@@ -119,10 +115,6 @@ class NodeBadgeExtension implements ComfyExtension {
   }
 
   nodeCreated(node: ComfyLGraphNode, app: ComfyApp) {
-    if (!app.vueAppReady) {
-      return
-    }
-
     node.badgePosition = BadgePosition.TopRight
     // @ts-expect-error Disable ComfyUI-Manager's badge drawing by setting badge_enabled to true. Remove this when ComfyUI-Manager's badge drawing is removed.
     node.badge_enabled = true

diff --git a/src/scripts/app.ts b/src/scripts/app.ts
@@ -1895,9 +1895,7 @@ export class ComfyApp {
 
     // Save current workflow automatically
     setInterval(() => {
-      const sortNodes =
-        this.vueAppReady &&
-        useSettingStore().get('Comfy.Workflow.SortNodeIdOnSave')
+      const sortNodes = useSettingStore().get('Comfy.Workflow.SortNodeIdOnSave')
       const workflow = JSON.stringify(this.graph.serialize({ sortNodes }))
       localStorage.setItem('workflow', workflow)
       if (api.clientId) {
@@ -2134,10 +2132,7 @@ export class ComfyApp {
   }
 
   showMissingNodesError(missingNodeTypes, hasAddedNodes = true) {
-    if (
-      this.vueAppReady &&
-      useSettingStore().get('Comfy.Workflow.ShowMissingNodesWarning')
-    ) {
+    if (useSettingStore().get('Comfy.Workflow.ShowMissingNodesWarning')) {
       showLoadWorkflowWarning({
         missingNodeTypes,
         hasAddedNodes,
@@ -2151,10 +2146,7 @@ export class ComfyApp {
   }
 
   showMissingModelsError(missingModels) {
-    if (
-      this.vueAppReady &&
-      useSettingStore().get('Comfy.Workflow.ShowMissingModelsWarning')
-    ) {
+    if (useSettingStore().get('Comfy.Workflow.ShowMissingModelsWarning')) {
       showMissingModelsWarning({
         missingModels,
         maximizable: true
@@ -2210,10 +2202,7 @@ export class ComfyApp {
       console.error(error)
     }
 
-    if (
-      this.vueAppReady &&
-      useSettingStore().get('Comfy.Validation.Workflows')
-    ) {
+    if (useSettingStore().get('Comfy.Validation.Workflows')) {
       // TODO: Show validation error in a dialog.
       const validatedGraphData = await validateComfyWorkflow(
         graphData,
@@ -2416,9 +2405,8 @@ export class ComfyApp {
       }
     }
 
-    const sortNodes =
-      this.vueAppReady &&
-      useSettingStore().get('Comfy.Workflow.SortNodeIdOnSave')
+    const sortNodes = useSettingStore().get('Comfy.Workflow.SortNodeIdOnSave')
+
     const workflow = graph.serialize({ sortNodes })
     const output = {}
     // Process nodes in order of execution

diff --git a/tests-ui/utils/index.ts b/tests-ui/utils/index.ts
@@ -1,4 +1,4 @@
-import { APIConfig, mockApi } from './setup'
+import { APIConfig, mockApi, mockSettingStore } from './setup'
 import { Ez, EzGraph, EzNameSpace } from './ezgraph'
 import lg from './litegraph'
 import fs from 'fs'
@@ -40,7 +40,11 @@ export async function start(config: StartConfig = {}): Promise<StartResult> {
   document.body.innerHTML = html.toString()
 
   mockApi(config)
+  mockSettingStore()
   const { app } = await import('../../src/scripts/app')
+  const { useSettingStore } = await import('../../src/stores/settingStore')
+  useSettingStore().addSettings(app.ui.settings)
+
   const { LiteGraph, LGraphCanvas } = await import('@comfyorg/litegraph')
   config.preSetup?.(app)
   const canvasEl = document.createElement('canvas')

diff --git a/tests-ui/utils/setup.ts b/tests-ui/utils/setup.ts
@@ -1,3 +1,5 @@
+import type { ComfySettingsDialog } from '@/scripts/ui/settings'
+import type { ComfyApp } from '@/scripts/app'
 import '../../src/scripts/api'
 
 const fs = require('fs')
@@ -97,3 +99,28 @@ export function mockApi(config: APIConfig = {}) {
     }
   }))
 }
+
+export const mockSettingStore = () => {
+  let app: ComfyApp | null = null
+
+  const mockedSettingStore = {
+    addSettings(settings: ComfySettingsDialog) {
+      app = settings.app
+    },
+
+    set(key: string, value: any) {
+      app?.ui.settings.setSettingValue(key, value)
+    },
+
+    get(key: string) {
+      return (
+        app?.ui.settings.getSettingValue(key) ??
+        app?.ui.settings.getSettingDefaultValue(key)
+      )
+    }
+  }
+
+  jest.mock('@/stores/settingStore', () => ({
+    useSettingStore: jest.fn(() => mockedSettingStore)
+  }))
+}