test: strengthen ECS migration coverage

Amp-Thread-ID: https://ampcode.com/threads/T-01a00769-0047-7184-b6ca-2bc56dcd5f71

Author: ampagent

src/lib/litegraph/src/__fixtures__/nodeHelpers.ts

@@ -47,6 +47,13 @@ export function createTestWidgetNode(graph: LGraph | Subgraph): LGraphNode {
       }
     }
     LiteGraph.registered_node_types[WIDGET_NODE_TYPE] = WidgetTestNode
+    onTestFinished(() => {
+      if (
+        LiteGraph.registered_node_types[WIDGET_NODE_TYPE] === WidgetTestNode
+      ) {
+        delete LiteGraph.registered_node_types[WIDGET_NODE_TYPE]
+      }
+    })
   }
   const node = LiteGraph.createNode(WIDGET_NODE_TYPE)
   if (!node) throw new Error('Failed to create widget node')

src/lib/litegraph/src/node/slotUtils.test.ts

@@ -5,21 +5,23 @@ import { beforeEach, describe, expect, it } from 'vitest'
 import type { INodeOutputSlot } from '@/lib/litegraph/src/interfaces'
 import type { IWidget } from '@/lib/litegraph/src/litegraph'
 import { LGraph, LGraphNode } from '@/lib/litegraph/src/litegraph'
+import { toLinkId } from '@/types/linkId'
 
 import { outputAsSerialisable } from './slotUtils'
 
 type OutputSlotParam = INodeOutputSlot & { widget?: IWidget }
 
-function createConnectedGraph(targetCount: number) {
+function createConnectedGraph(linkIds: number[]) {
   const graph = new LGraph()
   const source = new LGraphNode('Source')
   source.addOutput('out', 'number')
   graph.add(source)
 
-  for (let i = 0; i < targetCount; i++) {
+  for (const [i, linkId] of linkIds.entries()) {
     const target = new LGraphNode(`Target${i}`)
     target.addInput('in', 'number')
     graph.add(target)
+    graph.state.lastLinkId = toLinkId(linkId - 1)
     source.connect(0, target, 0)
   }
 
@@ -30,20 +32,22 @@ describe('outputAsSerialisable', () => {
   beforeEach(() => setActivePinia(createTestingPinia({ stubActions: false })))
 
   it('serialises the links leaving the slot, ascending by id', () => {
-    const { source } = createConnectedGraph(3)
+    const { source } = createConnectedGraph([10, 2])
 
     const serialised = outputAsSerialisable(
       source.outputs[0] as OutputSlotParam,
       source,
       0
     )
 
-    expect(serialised.links).toHaveLength(3)
-    expect(serialised.links).toEqual([...serialised.links!].sort())
+    expect(serialised.links).toEqual(
+      [...serialised.links!].sort((a, b) => a - b)
+    )
+    expect(serialised.links).toEqual([2, 10])
   })
 
   it('returns a snapshot unaffected by later graph changes', () => {
-    const { source } = createConnectedGraph(2)
+    const { source } = createConnectedGraph([1, 2])
 
     const serialised = outputAsSerialisable(
       source.outputs[0] as OutputSlotParam,
@@ -57,7 +61,7 @@ describe('outputAsSerialisable', () => {
   })
 
   it('serialises null when the slot has no links', () => {
-    const { source } = createConnectedGraph(0)
+    const { source } = createConnectedGraph([])
 
     const serialised = outputAsSerialisable(
       source.outputs[0] as OutputSlotParam,

src/lib/litegraph/src/subgraph/__fixtures__/subgraphHelpers.test.ts

@@ -66,6 +66,28 @@ describe('enableSubgraphNodeCreation', () => {
     dispose()
 
     expect(LiteGraph.createNode(subgraph.id)).toBeNull()
+
+    const laterSubgraph = rootGraph.createSubgraph(createTestSubgraphData())
+    expect(LiteGraph.createNode(laterSubgraph.id)).toBeNull()
+  })
+
+  it('only removes registrations owned by its invocation', () => {
+    const firstRoot = createTestRootGraph()
+    const secondRoot = createTestRootGraph()
+    const disposeFirst = enableSubgraphNodeCreation(firstRoot)
+    const disposeSecond = enableSubgraphNodeCreation(secondRoot)
+    const data = createTestSubgraphData()
+    firstRoot.createSubgraph(data)
+    const secondSubgraph = secondRoot.createSubgraph(data)
+
+    disposeFirst()
+
+    expect(LiteGraph.createNode(data.id)).toMatchObject({
+      subgraph: secondSubgraph
+    })
+
+    disposeSecond()
+    expect(LiteGraph.createNode(data.id)).toBeNull()
   })
 })
 

src/lib/litegraph/src/subgraph/__fixtures__/subgraphHelpers.ts

@@ -17,6 +17,7 @@ import {
   SUBGRAPH_INPUT_ID,
   SUBGRAPH_OUTPUT_ID
 } from '@/lib/litegraph/src/constants'
+import type { LGraphEventMap } from '@/lib/litegraph/src/infrastructure/LGraphEventMap'
 import type { SerializedNodeId } from '@/types/nodeId'
 import { toNodeId } from '@/types/nodeId'
 import {
@@ -60,22 +61,15 @@ function nextFixtureUuid(): UUID {
 export function resetSubgraphFixtureState(): void {
   fixtureUuidSequence = 1
   cleanupComplexPromotionFixtureNodeType()
-  disposeSubgraphNodeCreation()
-}
-
-const subgraphNodeTypesToDispose: string[] = []
-
-function disposeSubgraphNodeCreation(): void {
-  for (const type of subgraphNodeTypesToDispose) {
-    delete LiteGraph.registered_node_types[type]
-  }
-  subgraphNodeTypesToDispose.length = 0
 }
 
 export function enableSubgraphNodeCreation(rootGraph: LGraph): () => void {
-  rootGraph.events.addEventListener('subgraph-created', (e) => {
+  const registrations = new Map<string, typeof LGraphNode>()
+  const listener = (
+    e: CustomEvent<LGraphEventMap['subgraph-created']>
+  ): void => {
     const { subgraph } = e.detail
-    LiteGraph.registered_node_types[subgraph.id] = class extends SubgraphNode {
+    class TestSubgraphNode extends SubgraphNode {
       constructor() {
         super(rootGraph, subgraph, {
           id: -1,
@@ -90,9 +84,19 @@ export function enableSubgraphNodeCreation(rootGraph: LGraph): () => void {
         })
       }
     }
-    subgraphNodeTypesToDispose.push(subgraph.id)
-  })
-  return disposeSubgraphNodeCreation
+    LiteGraph.registered_node_types[subgraph.id] = TestSubgraphNode
+    registrations.set(subgraph.id, TestSubgraphNode)
+  }
+  rootGraph.events.addEventListener('subgraph-created', listener)
+
+  return () => {
+    rootGraph.events.removeEventListener('subgraph-created', listener)
+    for (const [type, constructor] of registrations) {
+      if (LiteGraph.registered_node_types[type] === constructor) {
+        delete LiteGraph.registered_node_types[type]
+      }
+    }
+  }
 }
 
 export function createTestRootGraph(id: UUID = nextFixtureUuid()): LGraph {

src/renderer/core/canvas/litegraph/notifyLayoutChanges.test.ts

@@ -84,6 +84,7 @@ describe('notifyLayoutChanges', () => {
 
     resize(context)
     await vi.waitFor(() => expect(onResize).toHaveBeenCalled())
+    expect([...onResize.mock.calls[0][0]]).toEqual([300, 200])
   })
 
   it('leaves onResize alone when a bounds batch only moves', async () => {

src/systems/badgeSystem.pricing.test.ts

@@ -4,27 +4,15 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'
 
 import { LGraph, LGraphNode } from '@/lib/litegraph/src/litegraph'
 import { useLinkStore } from '@/stores/linkStore'
-import {
-  graphScopeOf,
-  toOwningGraphId,
-  toRootGraphId
-} from '@/types/graphScopeId'
+import { toOwningGraphId, toRootGraphId } from '@/types/graphScopeId'
 import { toLinkId } from '@/types/linkId'
 import { toNodeId } from '@/types/nodeId'
 import type { UUID } from '@/utils/uuid'
 
 import { nodeBadges } from './badgeSystem'
 
-const getNodeDisplayPrice = vi.fn((node: LGraphNode) => {
-  const graph = node.graph
-  if (!graph) return '$disconnected'
-  const connected = node.inputs.some(
-    (input, index) =>
-      (input.name === 'image' || input.name?.startsWith('ref_images.')) &&
-      useLinkStore().isInputSlotConnected(graphScopeOf(graph), node.id, index)
-  )
-  return connected ? '$connected' : '$disconnected'
-})
+let displayPrice = '$disconnected'
+const getNodeDisplayPrice = vi.fn(() => displayPrice)
 
 vi.mock('@/composables/node/useNodePricing', () => {
   return {
@@ -59,6 +47,7 @@ function scopeOf(id: string) {
 describe('badge derivation pricing input connectivity', () => {
   beforeEach(() => {
     setActivePinia(createTestingPinia({ stubActions: false }))
+    displayPrice = '$disconnected'
     getNodeDisplayPrice.mockClear()
   })
 
@@ -89,9 +78,11 @@ describe('badge derivation pricing input connectivity', () => {
     const { node } = setup(['image', 'other'])
     expect(nodeBadges(node).at(-1)?.text).toBe('$disconnected')
 
+    displayPrice = '$connected'
     connect(0, 1)
     expect(nodeBadges(node).at(-1)?.text).toBe('$connected')
 
+    displayPrice = '$disconnected'
     const linkStore = useLinkStore()
     const topology = linkStore.getInputSlotLink(
       scopeOf(graphId),
@@ -107,6 +98,7 @@ describe('badge derivation pricing input connectivity', () => {
 
     expect(nodeBadges(node).at(-1)?.text).toBe('$disconnected')
 
+    displayPrice = '$connected'
     connect(1, 2)
     expect(nodeBadges(node).at(-1)?.text).toBe('$connected')
   })
@@ -117,6 +109,7 @@ describe('badge derivation pricing input connectivity', () => {
 
     expect(nodeBadges(node).at(-1)?.text).toBe('$disconnected')
 
+    displayPrice = '$connected'
     connect(1, 3)
     expect(nodeBadges(node).at(-1)?.text).toBe('$disconnected')
   })
@@ -127,6 +120,7 @@ describe('badge derivation pricing input connectivity', () => {
 
     const reloadedGraphId: UUID = 'graph-pricing-reloaded'
     node.graph!.rootGraph.id = reloadedGraphId
+    displayPrice = '$connected'
     connect(0, 4, reloadedGraphId)
 
     expect(nodeBadges(node).at(-1)?.text).toBe('$connected')