ComfyUI_frontend/src/components/graph/SelectionToolbox.test.ts

/* eslint-disable testing-library/no-container, testing-library/no-node-access */
import { createTestingPinia } from '@pinia/testing'
import { fireEvent, render } from '@testing-library/vue'
import { setActivePinia } from 'pinia'
import PrimeVue from 'primevue/config'
import { beforeEach, describe, expect, it, vi } from 'vitest'
import { createI18n } from 'vue-i18n'

import SelectionToolbox from '@/components/graph/SelectionToolbox.vue'
import type { LGraphNode } from '@/lib/litegraph/src/LGraphNode'
import { useCanvasStore } from '@/renderer/core/canvas/canvasStore'
import { useCanvasInteractions } from '@/renderer/core/canvas/useCanvasInteractions'
import { useExtensionService } from '@/services/extensionService'
import {
  createMockCanvas,
  createMockPositionable
} from '@/utils/__tests__/litegraphTestUtils'
import * as litegraphUtil from '@/utils/litegraphUtil'
import * as nodeFilterUtil from '@/utils/nodeFilterUtil'

function createMockExtensionService(): ReturnType<typeof useExtensionService> {
  return {
    extensionCommands: { value: new Map() },
    loadExtensions: vi.fn(),
    registerExtension: vi.fn(),
    invokeExtensions: vi.fn(() => []),
    invokeExtensionsAsync: vi.fn()
  } as Partial<ReturnType<typeof useExtensionService>> as ReturnType<
    typeof useExtensionService
  >
}

const { settingGetMock } = vi.hoisted(() => ({
  settingGetMock: vi.fn()
}))

const defaultSettingValues: Record<string, unknown> = {
  'Comfy.UseNewMenu': 'Top',
  'Comfy.NodeLibrary.NewDesign': true,
  'Comfy.Load3D.3DViewerEnable': true
}

function mockSettingValues(overrides: Record<string, unknown> = {}) {
  const settingValues = {
    ...defaultSettingValues,
    ...overrides
  }
  settingGetMock.mockImplementation(
    (key: string): unknown => settingValues[key] ?? null
  )
}

// Mock the composables and services
vi.mock('@/renderer/core/canvas/useCanvasInteractions', () => ({
  useCanvasInteractions: vi.fn(() => ({
    handleWheel: vi.fn()
  }))
}))

vi.mock('@/composables/canvas/useSelectionToolboxPosition', () => ({
  useSelectionToolboxPosition: vi.fn(() => ({
    visible: { value: true }
  })),
  resetMoreOptionsState: vi.fn()
}))

vi.mock('@/composables/element/useRetriggerableAnimation', () => ({
  useRetriggerableAnimation: vi.fn(() => ({
    shouldAnimate: { value: false }
  }))
}))

vi.mock('@/renderer/extensions/minimap/composables/useMinimap', () => ({
  useMinimap: vi.fn(() => ({
    containerStyles: {
      value: {
        backgroundColor: '#ffffff'
      }
    }
  }))
}))

vi.mock('@/services/extensionService', () => ({
  useExtensionService: vi.fn(() => ({
    extensionCommands: { value: new Map() },
    invokeExtensions: vi.fn(() => [])
  }))
}))

vi.mock('@/utils/litegraphUtil', () => ({
  isLGraphNode: vi.fn(() => true),
  isImageNode: vi.fn(() => false),
  isLoad3dNode: vi.fn(() => false)
}))

vi.mock('@/utils/nodeFilterUtil', () => ({
  isOutputNode: vi.fn(() => false),
  filterOutputNodes: vi.fn((nodes) => nodes.filter(() => false))
}))

vi.mock('@/platform/settings/settingStore', () => ({
  useSettingStore: () => ({
    get: settingGetMock
  })
}))

vi.mock('@/stores/commandStore', () => ({
  useCommandStore: () => ({
    getCommand: vi.fn(() => ({ id: 'test-command', title: 'Test Command' }))
  })
}))

let nodeDefMock = {
  type: 'TestNode',
  title: 'Test Node'
} as unknown

vi.mock('@/stores/nodeDefStore', () => ({
  useNodeDefStore: () => ({
    fromLGraphNode: vi.fn(() => nodeDefMock)
  })
}))

describe('SelectionToolbox', () => {
  let canvasStore: ReturnType<typeof useCanvasStore>

  const i18n = createI18n({
    legacy: false,
    locale: 'en',
    messages: {
      en: {
        g: {
          info: 'Node Info',
          bookmark: 'Save to Library',
          frameNodes: 'Frame Nodes',
          moreOptions: 'More Options',
          refreshNode: 'Refresh Node'
        }
      }
    }
  })

  const mockProvide = {
    isVisible: { value: true },
    selectedItems: []
  }

  beforeEach(() => {
    setActivePinia(createTestingPinia({ createSpy: vi.fn, stubActions: false }))
    canvasStore = useCanvasStore()
    nodeDefMock = {
      type: 'TestNode',
      title: 'Test Node'
    } as unknown

    // Mock the canvas to avoid "getCanvas: canvas is null" errors
    canvasStore.canvas = createMockCanvas()

    vi.resetAllMocks()
    mockSettingValues()
  })

  function renderComponent(props = {}): { container: Element } {
    const { container } = render(SelectionToolbox, {
      props,
      global: {
        plugins: [i18n, PrimeVue],
        provide: {
          [Symbol.for('SelectionOverlay')]: mockProvide
        },
        stubs: {
          Panel: {
            template:
              '<div class="panel selection-toolbox absolute left-1/2 rounded-lg"><slot /></div>',
            props: ['pt', 'style', 'class']
          },
          NodeContextMenu: { template: '<div class="node-context-menu" />' },
          InfoButton: { template: '<div class="info-button" />' },
          ColorPickerButton: {
            template:
              '<button data-testid="color-picker-button" class="color-picker-button" />'
          },
          FrameNodes: { template: '<div class="frame-nodes" />' },
          PublishButton: {
            template:
              '<button data-testid="add-to-library" class="bookmark-button" />'
          },
          BypassButton: {
            template:
              '<button data-testid="bypass-button" class="bypass-button" />'
          },
          PinButton: { template: '<div class="pin-button" />' },
          Load3DViewerButton: {
            template: '<div class="load-3d-viewer-button" />'
          },
          MaskEditorButton: {
            template: '<div class="mask-editor-button" />'
          },
          DeleteButton: {
            template:
              '<button data-testid="delete-button" class="delete-button" />'
          },
          RefreshSelectionButton: {
            template: '<div class="refresh-button" />'
          },
          ExecuteButton: { template: '<div class="execute-button" />' },
          ConvertToSubgraphButton: {
            template:
              '<button data-testid="convert-to-subgraph-button" class="convert-to-subgraph-button" />'
          },
          ExtensionCommandButton: {
            template: '<div class="extension-command-button" />'
          },
          MoreOptions: {
            template:
              '<button data-testid="more-options-button" class="more-options" />'
          },
          VerticalDivider: { template: '<div class="vertical-divider" />' }
        }
      }
    })
    return { container }
  }

  describe('Button Visibility Logic', () => {
    beforeEach(() => {
      const mockExtensionService = vi.mocked(useExtensionService)
      mockExtensionService.mockReturnValue(createMockExtensionService())
    })

    it('should show info button only for single selections', () => {
      // Single node selection
      canvasStore.selectedItems = [createMockPositionable()]
      const { container } = renderComponent()
      expect(container.querySelector('.info-button')).toBeTruthy()

      // Multiple node selection - render in separate test scope
      canvasStore.selectedItems = [
        createMockPositionable(),
        createMockPositionable()
      ]
      const { container: container2 } = renderComponent()
      expect(container2.querySelector('.info-button')).toBeFalsy()
    })

    it('should not show info button when node definition is not found', () => {
      canvasStore.selectedItems = [createMockPositionable()]
      nodeDefMock = null
      const { container } = renderComponent()
      expect(container.querySelector('.info-button')).toBeFalsy()
    })

    it('should not show info button when legacy menu uses the new node library', () => {
      mockSettingValues({
        'Comfy.UseNewMenu': 'Disabled',
        'Comfy.NodeLibrary.NewDesign': true
      })
      canvasStore.selectedItems = [createMockPositionable()]

      const { container } = renderComponent()

      expect(container.querySelector('.info-button')).toBeFalsy()
    })

    it('should not show info button when legacy menu uses the legacy node library', () => {
      mockSettingValues({
        'Comfy.UseNewMenu': 'Disabled',
        'Comfy.NodeLibrary.NewDesign': false
      })
      canvasStore.selectedItems = [createMockPositionable()]

      const { container } = renderComponent()

      expect(container.querySelector('.info-button')).toBeFalsy()
    })

    it('should show info button when new menu uses the legacy node library', () => {
      mockSettingValues({
        'Comfy.UseNewMenu': 'Top',
        'Comfy.NodeLibrary.NewDesign': false
      })
      canvasStore.selectedItems = [createMockPositionable()]

      const { container } = renderComponent()

      expect(container.querySelector('.info-button')).toBeTruthy()
    })

    it('should show color picker for all selections', () => {
      // Single node selection
      canvasStore.selectedItems = [createMockPositionable()]
      const { container } = renderComponent()
      expect(
        container.querySelector('[data-testid="color-picker-button"]')
      ).toBeTruthy()

      // Multiple node selection
      canvasStore.selectedItems = [
        createMockPositionable(),
        createMockPositionable()
      ]
      const { container: container2 } = renderComponent()
      expect(
        container2.querySelector('[data-testid="color-picker-button"]')
      ).toBeTruthy()
    })

    it('should show frame nodes only for multiple selections', () => {
      // Single node selection
      canvasStore.selectedItems = [createMockPositionable()]
      const { container } = renderComponent()
      expect(container.querySelector('.frame-nodes')).toBeFalsy()

      // Multiple node selection
      canvasStore.selectedItems = [
        createMockPositionable(),
        createMockPositionable()
      ]
      const { container: container2 } = renderComponent()
      expect(container2.querySelector('.frame-nodes')).toBeTruthy()
    })

    it('should show bypass button for appropriate selections', () => {
      // Single node selection
      canvasStore.selectedItems = [createMockPositionable()]
      const { container } = renderComponent()
      expect(
        container.querySelector('[data-testid="bypass-button"]')
      ).toBeTruthy()

      // Multiple node selection
      canvasStore.selectedItems = [
        createMockPositionable(),
        createMockPositionable()
      ]
      const { container: container2 } = renderComponent()
      expect(
        container2.querySelector('[data-testid="bypass-button"]')
      ).toBeTruthy()
    })

    it('should show common buttons for all selections', () => {
      canvasStore.selectedItems = [createMockPositionable()]
      const { container } = renderComponent()

      expect(
        container.querySelector('[data-testid="delete-button"]')
      ).toBeTruthy()
      expect(
        container.querySelector('[data-testid="convert-to-subgraph-button"]')
      ).toBeTruthy()
      expect(
        container.querySelector('[data-testid="more-options-button"]')
      ).toBeTruthy()
    })

    it('should show mask editor only for single image nodes', () => {
      const isImageNodeSpy = vi.spyOn(litegraphUtil, 'isImageNode')

      // Single image node
      isImageNodeSpy.mockReturnValue(true)
      canvasStore.selectedItems = [createMockPositionable()]
      const { container } = renderComponent()
      expect(container.querySelector('.mask-editor-button')).toBeTruthy()

      // Single non-image node
      isImageNodeSpy.mockReturnValue(false)
      canvasStore.selectedItems = [createMockPositionable()]
      const { container: container2 } = renderComponent()
      expect(container2.querySelector('.mask-editor-button')).toBeFalsy()
    })

    it('should show Color picker button only for single Load3D nodes', () => {
      const isLoad3dNodeSpy = vi.spyOn(litegraphUtil, 'isLoad3dNode')

      // Single Load3D node
      isLoad3dNodeSpy.mockReturnValue(true)
      canvasStore.selectedItems = [createMockPositionable()]
      const { container } = renderComponent()
      expect(container.querySelector('.load-3d-viewer-button')).toBeTruthy()

      // Single non-Load3D node
      isLoad3dNodeSpy.mockReturnValue(false)
      canvasStore.selectedItems = [createMockPositionable()]
      const { container: container2 } = renderComponent()
      expect(container2.querySelector('.load-3d-viewer-button')).toBeFalsy()
    })

    it('should show ExecuteButton only when output nodes are selected', () => {
      const isOutputNodeSpy = vi.spyOn(nodeFilterUtil, 'isOutputNode')
      const filterOutputNodesSpy = vi.spyOn(nodeFilterUtil, 'filterOutputNodes')

      // With output node selected
      isOutputNodeSpy.mockReturnValue(true)
      filterOutputNodesSpy.mockReturnValue([
        { type: 'SaveImage' }
      ] as LGraphNode[])
      canvasStore.selectedItems = [createMockPositionable()]
      const { container } = renderComponent()
      expect(container.querySelector('.execute-button')).toBeTruthy()

      // Without output node selected
      isOutputNodeSpy.mockReturnValue(false)
      filterOutputNodesSpy.mockReturnValue([])
      canvasStore.selectedItems = [createMockPositionable()]
      const { container: container2 } = renderComponent()
      expect(container2.querySelector('.execute-button')).toBeFalsy()

      // No selection at all
      canvasStore.selectedItems = []
      const { container: container3 } = renderComponent()
      expect(container3.querySelector('.execute-button')).toBeFalsy()
    })
  })

  describe('Divider Visibility Logic', () => {
    it('should show dividers between button groups when both groups have buttons', () => {
      // Setup single node to show info + other buttons
      canvasStore.selectedItems = [createMockPositionable()]
      const { container } = renderComponent()

      const dividers = container.querySelectorAll('.vertical-divider')
      expect(dividers.length).toBeGreaterThan(0)
    })

    it('should not show dividers when adjacent groups are empty', () => {
      // No selection should show minimal buttons and dividers
      canvasStore.selectedItems = []
      const { container } = renderComponent()

      expect(
        container.querySelector('[data-testid="more-options-button"]')
      ).toBeTruthy()
    })
  })

  describe('Extension Commands', () => {
    it('should render extension command buttons when available', () => {
      const mockExtensionService = vi.mocked(useExtensionService)
      mockExtensionService.mockReturnValue({
        extensionCommands: {
          value: new Map([
            ['test-command', { id: 'test-command', title: 'Test Command' }]
          ])
        },
        loadExtensions: vi.fn(),
        registerExtension: vi.fn(),
        invokeExtensions: vi.fn(() => ['test-command']),
        invokeExtensionsAsync: vi.fn()
      } as ReturnType<typeof useExtensionService>)

      canvasStore.selectedItems = [createMockPositionable()]
      const { container } = renderComponent()

      expect(container.querySelector('.extension-command-button')).toBeTruthy()
    })

    it('should not render extension commands when none available', () => {
      const mockExtensionService = vi.mocked(useExtensionService)
      mockExtensionService.mockReturnValue(createMockExtensionService())

      canvasStore.selectedItems = [createMockPositionable()]
      const { container } = renderComponent()

      expect(container.querySelector('.extension-command-button')).toBeFalsy()
    })
  })

  describe('Event Handling', () => {
    it('should handle wheel events', async () => {
      const mockCanvasInteractions = vi.mocked(useCanvasInteractions)
      const forwardEventToCanvasSpy = vi.fn()
      mockCanvasInteractions.mockReturnValue({
        handleWheel: vi.fn(),
        handlePointer: vi.fn(),
        forwardEventToCanvas: forwardEventToCanvasSpy,
        shouldHandleNodePointerEvents: { value: true } as ReturnType<
          typeof useCanvasInteractions
        >['shouldHandleNodePointerEvents']
      } as ReturnType<typeof useCanvasInteractions>)

      const mockExtensionService = vi.mocked(useExtensionService)
      mockExtensionService.mockReturnValue(createMockExtensionService())

      canvasStore.selectedItems = [createMockPositionable()]
      const { container } = renderComponent()

      const panel = container.querySelector('.panel')
      expect(panel).toBeTruthy()
      await fireEvent.wheel(panel!)

      expect(forwardEventToCanvasSpy).toHaveBeenCalled()
    })
  })

  describe('No Selection State', () => {
    beforeEach(() => {
      const mockExtensionService = vi.mocked(useExtensionService)
      mockExtensionService.mockReturnValue(createMockExtensionService())
    })

    it('should hide most buttons when no items selected', () => {
      canvasStore.selectedItems = []
      const { container } = renderComponent()

      expect(container.querySelector('.info-button')).toBeFalsy()
      expect(container.querySelector('.color-picker-button')).toBeFalsy()
      expect(container.querySelector('.frame-nodes')).toBeFalsy()
      expect(container.querySelector('.bookmark-button')).toBeFalsy()
    })
  })
})