test: extract shared builder helper and expand appModeStore coverage

Address review feedback: move saveCloseAndReopenInBuilder to builderTestUtils, reuse it in builderReorder.spec.ts, and add unit tests for enterBuilder/exitBuilder pruning, removeSelectedInput, updateInputConfig merge, loadSelections edge cases, and pruneLinearData branches.

Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent)

Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>

browser_tests/fixtures/utils/builderTestUtils.ts

@@ -122,3 +122,19 @@ export async function saveAndReopenInAppMode(
 
   await comfyPage.appMode.toggleAppMode()
 }
+
+export async function saveCloseAndReopenInBuilder(
+  comfyPage: ComfyPage,
+  appMode: AppModeHelper,
+  workflowName: string
+) {
+  await appMode.steps.goToPreview()
+  await builderSaveAs(appMode, workflowName)
+  await appMode.saveAs.closeButton.click()
+  await expect(appMode.saveAs.successDialog).toBeHidden()
+
+  await appMode.footer.exitBuilder()
+  await openWorkflowFromSidebar(comfyPage, workflowName)
+  await appMode.enterBuilder()
+  await appMode.steps.goToInputs()
+}

browser_tests/tests/appModeInputPersistence.spec.ts

@@ -2,32 +2,13 @@ import {
   comfyPageFixture as test,
   comfyExpect as expect
 } from '@e2e/fixtures/ComfyPage'
-import type { AppModeHelper } from '@e2e/fixtures/helpers/AppModeHelper'
-import type { ComfyPage } from '@e2e/fixtures/ComfyPage'
 import {
-  builderSaveAs,
-  openWorkflowFromSidebar,
+  saveCloseAndReopenInBuilder,
   setupBuilder
 } from '@e2e/helpers/builderTestUtils'
 
 const WIDGETS = ['seed', 'steps', 'cfg']
 
-async function saveCloseAndReopenInBuilder(
-  comfyPage: ComfyPage,
-  appMode: AppModeHelper,
-  workflowName: string
-) {
-  await appMode.steps.goToPreview()
-  await builderSaveAs(appMode, workflowName)
-  await appMode.saveAs.closeButton.click()
-  await expect(appMode.saveAs.successDialog).toBeHidden()
-
-  await appMode.footer.exitBuilder()
-  await openWorkflowFromSidebar(comfyPage, workflowName)
-  await appMode.enterBuilder()
-  await appMode.steps.goToInputs()
-}
-
 test.describe(
   'App builder input persistence after reload',
   { tag: '@ui' },

browser_tests/tests/builderReorder.spec.ts

@@ -5,26 +5,18 @@ import {
 import type { AppModeHelper } from '@e2e/fixtures/helpers/AppModeHelper'
 import type { ComfyPage } from '@e2e/fixtures/ComfyPage'
 import {
-  builderSaveAs,
-  openWorkflowFromSidebar,
+  saveCloseAndReopenInBuilder,
   setupBuilder
 } from '@e2e/fixtures/utils/builderTestUtils'
 
 const WIDGETS = ['seed', 'steps', 'cfg']
 
-/** Save as app, close it by loading default, reopen from sidebar, enter app mode. */
 async function saveCloseAndReopenAsApp(
   comfyPage: ComfyPage,
   appMode: AppModeHelper,
   workflowName: string
 ) {
-  await appMode.steps.goToPreview()
-  await builderSaveAs(appMode, workflowName)
-  await appMode.saveAs.closeButton.click()
-  await expect(appMode.saveAs.successDialog).toBeHidden()
-
-  await appMode.footer.exitBuilder()
-  await openWorkflowFromSidebar(comfyPage, workflowName)
+  await saveCloseAndReopenInBuilder(comfyPage, appMode, workflowName)
   await appMode.toggleAppMode()
 }
 

src/stores/appModeStore.test.ts

@@ -5,7 +5,12 @@ import { nextTick } from 'vue'
 import { beforeEach, describe, expect, it, vi } from 'vitest'
 
 import type { LGraphNode, NodeId } from '@/lib/litegraph/src/LGraphNode'
-import type { LoadedComfyWorkflow } from '@/platform/workflow/management/stores/comfyWorkflow'
+import type { IBaseWidget } from '@/lib/litegraph/src/types/widgets'
+import type {
+  InputWidgetConfig,
+  LinearInput,
+  LoadedComfyWorkflow
+} from '@/platform/workflow/management/stores/comfyWorkflow'
 import { ComfyWorkflow as ComfyWorkflowClass } from '@/platform/workflow/management/stores/comfyWorkflow'
 import { useWorkflowStore } from '@/platform/workflow/management/stores/workflowStore'
 import { app } from '@/scripts/app'
@@ -100,6 +105,29 @@ function createBuilderWorkflowWithOutputs(
   return workflow
 }
 
+function createWorkflowWithLinearData(
+  activeMode: string,
+  inputs: LinearInput[],
+  outputs: NodeId[]
+): LoadedComfyWorkflow {
+  const workflow = createBuilderWorkflow(activeMode)
+  workflow.changeTracker = createMockChangeTracker(
+    fromPartial<Partial<ChangeTracker>>({
+      activeState: {
+        last_node_id: 0,
+        last_link_id: 0,
+        nodes: [],
+        links: [],
+        groups: [],
+        config: {},
+        version: 0.4,
+        extra: { linearData: fromAny({ inputs, outputs }) }
+      }
+    })
+  )
+  return workflow
+}
+
 describe('appModeStore', () => {
   let workflowStore: ReturnType<typeof useWorkflowStore>
   let store: ReturnType<typeof useAppModeStore>
@@ -164,6 +192,28 @@ describe('appModeStore', () => {
       )
       expect(workflowStore.activeWorkflow!.activeMode).toBe('graph')
     })
+
+    it('prunes selections from workflow state on entry', () => {
+      const node1 = { id: 1 }
+      mockResolveNode.mockImplementation((id) =>
+        id == 1 ? fromAny<LGraphNode, unknown>(node1) : undefined
+      )
+      workflowStore.activeWorkflow = createWorkflowWithLinearData(
+        'graph',
+        [
+          [1, 'seed'],
+          [99, 'steps']
+        ],
+        [1, 99]
+      )
+      store.selectedInputs = [[42, 'prompt']]
+      store.selectedOutputs = [42]
+
+      store.enterBuilder()
+
+      expect(store.selectedInputs).toEqual([[1, 'seed']])
+      expect(store.selectedOutputs).toEqual([1])
+    })
   })
 
   describe('empty workflow dialog callbacks', () => {
@@ -203,33 +253,36 @@ describe('appModeStore', () => {
     })
   })
 
+  describe('exitBuilder', () => {
+    it('prunes selections from workflow state on exit', () => {
+      const node1 = { id: 1 }
+      mockResolveNode.mockImplementation((id) =>
+        id == 1 ? fromAny<LGraphNode, unknown>(node1) : undefined
+      )
+      workflowStore.activeWorkflow = createWorkflowWithLinearData(
+        'builder:inputs',
+        [
+          [1, 'seed'],
+          [99, 'steps']
+        ],
+        [1, 99]
+      )
+      store.selectedInputs = [[42, 'prompt']]
+      store.selectedOutputs = [42]
+
+      store.exitBuilder()
+
+      expect(store.selectedInputs).toEqual([[1, 'seed']])
+      expect(store.selectedOutputs).toEqual([1])
+      expect(workflowStore.activeWorkflow!.activeMode).toBe('graph')
+    })
+  })
+
   describe('loadSelections pruning', () => {
     function mockNode(id: number) {
       return { id }
     }
 
-    function workflowWithLinearData(
-      inputs: [number, string][],
-      outputs: number[]
-    ) {
-      const workflow = createBuilderWorkflow('app')
-      workflow.changeTracker = createMockChangeTracker(
-        fromPartial<Partial<ChangeTracker>>({
-          activeState: {
-            last_node_id: 0,
-            last_link_id: 0,
-            nodes: [],
-            links: [],
-            groups: [],
-            config: {},
-            version: 0.4,
-            extra: { linearData: { inputs, outputs } }
-          }
-        })
-      )
-      return workflow
-    }
-
     it('removes inputs referencing deleted nodes on load', async () => {
       const node1 = mockNode(1)
       mockResolveNode.mockImplementation((id) =>
@@ -295,7 +348,11 @@ describe('appModeStore', () => {
       // Initially nodes are not resolvable — pruning removes them
       mockResolveNode.mockReturnValue(undefined)
       const inputs: [number, string][] = [[1, 'seed']]
-      workflowStore.activeWorkflow = workflowWithLinearData(inputs, [1])
+      workflowStore.activeWorkflow = createWorkflowWithLinearData(
+        'app',
+        inputs,
+        [1]
+      )
       store.loadSelections({ inputs })
       await nextTick()
 
@@ -325,6 +382,49 @@ describe('appModeStore', () => {
     })
   })
 
+  describe('loadSelections edge cases', () => {
+    it('handles empty or undefined data', () => {
+      store.loadSelections(undefined)
+
+      expect(store.selectedInputs).toEqual([])
+      expect(store.selectedOutputs).toEqual([])
+
+      store.loadSelections({})
+
+      expect(store.selectedInputs).toEqual([])
+      expect(store.selectedOutputs).toEqual([])
+    })
+  })
+
+  describe('pruneLinearData', () => {
+    it('returns empty selections for undefined data', () => {
+      expect(store.pruneLinearData(undefined)).toEqual({
+        inputs: [],
+        outputs: []
+      })
+    })
+
+    it('does not prune when rootGraph is empty', () => {
+      const originalRootGraph = app.rootGraph
+      Object.defineProperty(app, 'rootGraph', { value: null, writable: true })
+
+      expect(
+        store.pruneLinearData({
+          inputs: [[1, 'seed']],
+          outputs: [1]
+        })
+      ).toEqual({
+        inputs: [[1, 'seed']],
+        outputs: [1]
+      })
+
+      Object.defineProperty(app, 'rootGraph', {
+        value: originalRootGraph,
+        writable: true
+      })
+    })
+  })
+
   describe('pruneLinearData during graph loading', () => {
     it('preserves all entries when ChangeTracker.isLoadingGraph is true', () => {
       ChangeTracker.isLoadingGraph = true
@@ -513,6 +613,18 @@ describe('appModeStore', () => {
       expect(store.selectedInputs[0][2]).toEqual({ height: 200 })
     })
 
+    it('merges existing config with new values', () => {
+      const existingConfig: InputWidgetConfig & { width: number } = {
+        height: 120,
+        width: 240
+      }
+      store.selectedInputs.push([1, 'prompt', existingConfig])
+
+      store.updateInputConfig(1 as NodeId, 'prompt', { height: 300 })
+
+      expect(store.selectedInputs[0][2]).toEqual({ height: 300, width: 240 })
+    })
+
     it('triggers linearData sync watcher', async () => {
       workflowStore.activeWorkflow = createBuilderWorkflow()
       store.selectedInputs.push([42, 'prompt'])
@@ -528,6 +640,17 @@ describe('appModeStore', () => {
     })
   })
 
+  describe('removeSelectedInput', () => {
+    it('removes the matching input entry only', () => {
+      store.selectedInputs.push([1, 'prompt'])
+      store.selectedInputs.push([2, 'steps'])
+
+      store.removeSelectedInput({ name: 'steps' } as IBaseWidget, { id: 2 })
+
+      expect(store.selectedInputs).toEqual([[1, 'prompt']])
+    })
+  })
+
   describe('autoEnableVueNodes', () => {
     it('enables Vue nodes when entering select mode with them disabled', async () => {
       mockSettings.store['Comfy.VueNodes.Enabled'] = false