Fix node badge on unknown color palette (#1487)

* Fix node badge on unknown color palette * Update test expectations [skip ci] --------- Co-authored-by: github-actions <[email protected]>
Comfy-Org · Nov 9, 2024 · 4617e0f · 4617e0f
1 parent f8ec87d
commit 4617e0f
Show file tree

Hide file tree

Showing 6 changed files with 26 additions and 5 deletions.
diff --git a/browser_tests/colorPalette.spec.ts b/browser_tests/colorPalette.spec.ts
@@ -211,7 +211,7 @@ test.describe('Node Color Adjustments', () => {
       await comfyPage.setSetting('Comfy.ColorPalette', 'light')
       await comfyPage.setSetting('Comfy.Node.Opacity', 0.3)
       const node = await comfyPage.getFirstNodeRef()
-      await node.clickContextMenuOption('Colors')
+      await node?.clickContextMenuOption('Colors')
     })
 
     test('should persist color adjustments when changing custom node colors', async ({

diff --git a/browser_tests/nodeBadge.spec.ts b/browser_tests/nodeBadge.spec.ts
@@ -74,3 +74,21 @@ test.describe('Node source badge', () => {
     })
   })
 })
+
+test.describe('Node badge color', () => {
+  test('Can show node badge with unknown color palette', async ({
+    comfyPage
+  }) => {
+    await comfyPage.setSetting(
+      'Comfy.NodeBadge.NodeIdBadgeMode',
+      NodeBadgeMode.ShowAll
+    )
+    await comfyPage.setSetting('Comfy.ColorPalette', 'unknown')
+    await comfyPage.nextFrame()
+    // Click empty space to trigger canvas re-render.
+    await comfyPage.clickEmptySpace()
+    await expect(comfyPage.canvas).toHaveScreenshot(
+      'node-badge-unknown-color-palette.png'
+    )
+  })
+})
diff --git a/...eBadge.spec.ts-snapshots/node-badge-unknown-color-palette-chromium-2x-linux.png b/...eBadge.spec.ts-snapshots/node-badge-unknown-color-palette-chromium-2x-linux.png
diff --git a/...nodeBadge.spec.ts-snapshots/node-badge-unknown-color-palette-chromium-linux.png b/...nodeBadge.spec.ts-snapshots/node-badge-unknown-color-palette-chromium-linux.png
diff --git a/src/components/graph/NodeBadge.vue b/src/components/graph/NodeBadge.vue
@@ -18,6 +18,7 @@ import { LGraphBadge } from '@comfyorg/litegraph'
 import _ from 'lodash'
 import { NodeBadgeMode } from '@/types/nodeSource'
 import { ComfyNodeDefImpl, useNodeDefStore } from '@/stores/nodeDefStore'
+import type { Palette } from '@/types/colorPalette'
 
 const settingStore = useSettingStore()
 const nodeSourceBadgeMode = computed(
@@ -35,7 +36,7 @@ watch([nodeSourceBadgeMode, nodeIdBadgeMode, nodeLifeCycleBadgeMode], () => {
   app.graph?.setDirtyCanvas(true, true)
 })
 
-const colorPalette = computed(() =>
+const colorPalette = computed<Palette | undefined>(() =>
   getColorPalette(settingStore.get('Comfy.ColorPalette'))
 )
 
@@ -78,10 +79,10 @@ onMounted(() => {
             }
           ),
           fgColor:
-            colorPalette.value.colors.litegraph_base?.BADGE_FG_COLOR ||
+            colorPalette.value?.colors?.litegraph_base?.BADGE_FG_COLOR ||
             defaultColorPalette.colors.litegraph_base.BADGE_FG_COLOR,
           bgColor:
-            colorPalette.value.colors.litegraph_base?.BADGE_BG_COLOR ||
+            colorPalette.value?.colors?.litegraph_base?.BADGE_BG_COLOR ||
             defaultColorPalette.colors.litegraph_base.BADGE_BG_COLOR
         })
       })

diff --git a/src/extensions/core/colorPalette.ts b/src/extensions/core/colorPalette.ts
@@ -452,7 +452,9 @@ const setCustomColorPalettes = (customColorPalettes: ColorPalettes) => {
 }
 
 export const defaultColorPalette = colorPalettes[defaultColorPaletteId]
-export const getColorPalette = (colorPaletteId?) => {
+export const getColorPalette = (
+  colorPaletteId?: string
+): Palette | undefined => {
   if (!colorPaletteId) {
     colorPaletteId = app.ui.settings.getSettingValue(id, defaultColorPaletteId)
   }