ComfyUI_frontend/src/composables/usePaste.test.ts

import { beforeEach, describe, expect, it, vi } from 'vitest'
import type {
  LGraphCanvas,
  LGraph,
  LGraphGroup,
  LGraphNode
} from '@/lib/litegraph/src/litegraph'
import { LiteGraph } from '@/lib/litegraph/src/litegraph'
import { app } from '@/scripts/app'
import { createNode, isImageNode } from '@/utils/litegraphUtil'
import {
  cloneDataTransfer,
  pasteImageNode,
  pasteImageNodes,
  usePaste
} from './usePaste'

function createMockNode() {
  return {
    pos: [0, 0],
    pasteFile: vi.fn(),
    pasteFiles: vi.fn()
  }
}

function createImageFile(
  name: string = 'test.png',
  type: string = 'image/png'
): File {
  return new File([''], name, { type })
}

function createAudioFile(
  name: string = 'test.mp3',
  type: string = 'audio/mpeg'
): File {
  return new File([''], name, { type })
}

function createDataTransfer(files: File[] = []): DataTransfer {
  const dataTransfer = new DataTransfer()
  files.forEach((file) => dataTransfer.items.add(file))
  return dataTransfer
}

const mockCanvas = {
  current_node: null as LGraphNode | null,
  graph: {
    add: vi.fn(),
    change: vi.fn()
  } as Partial<LGraph> as LGraph,
  graph_mouse: [100, 200],
  pasteFromClipboard: vi.fn(),
  _deserializeItems: vi.fn()
} as Partial<LGraphCanvas> as LGraphCanvas

const mockCanvasStore = {
  canvas: mockCanvas,
  getCanvas: vi.fn(() => mockCanvas)
}

const mockWorkspaceStore = {
  shiftDown: false
}

vi.mock('@vueuse/core', () => ({
  useEventListener: vi.fn((target, event, handler) => {
    target.addEventListener(event, handler)
    return () => target.removeEventListener(event, handler)
  })
}))

vi.mock('@/renderer/core/canvas/canvasStore', () => ({
  useCanvasStore: () => mockCanvasStore
}))

vi.mock('@/stores/workspaceStore', () => ({
  useWorkspaceStore: () => mockWorkspaceStore
}))

vi.mock('@/scripts/app', () => ({
  app: {
    loadGraphData: vi.fn()
  }
}))

vi.mock('@/lib/litegraph/src/litegraph', () => ({
  LiteGraph: {
    createNode: vi.fn()
  }
}))

vi.mock('@/utils/litegraphUtil', () => ({
  createNode: vi.fn(),
  isAudioNode: vi.fn(),
  isImageNode: vi.fn(),
  isVideoNode: vi.fn()
}))

vi.mock('@/workbench/eventHelpers', () => ({
  shouldIgnoreCopyPaste: vi.fn()
}))

describe('pasteImageNode', () => {
  beforeEach(() => {
    vi.clearAllMocks()
    vi.mocked(mockCanvas.graph!.add).mockImplementation(
      (node: LGraphNode | LGraphGroup | null) => node as LGraphNode
    )
  })

  it('should create new LoadImage node when no image node provided', async () => {
    const mockNode = createMockNode()
    vi.mocked(createNode).mockResolvedValue(mockNode as unknown as LGraphNode)

    const file = createImageFile()
    const dataTransfer = createDataTransfer([file])

    await pasteImageNode(
      mockCanvas as unknown as LGraphCanvas,
      dataTransfer.items
    )

    expect(createNode).toHaveBeenCalledWith(mockCanvas, 'LoadImage')
    expect(mockNode.pasteFile).toHaveBeenCalledWith(file)
  })

  it('should use existing image node when provided', async () => {
    const mockNode = createMockNode()
    const file = createImageFile()
    const dataTransfer = createDataTransfer([file])

    await pasteImageNode(
      mockCanvas as unknown as LGraphCanvas,
      dataTransfer.items,
      mockNode as unknown as LGraphNode
    )

    expect(mockNode.pasteFile).toHaveBeenCalledWith(file)
    expect(mockNode.pasteFiles).toHaveBeenCalledWith([file])
  })

  it('should handle multiple image files', async () => {
    const mockNode = createMockNode()
    const file1 = createImageFile('test1.png')
    const file2 = createImageFile('test2.jpg', 'image/jpeg')
    const dataTransfer = createDataTransfer([file1, file2])

    await pasteImageNode(
      mockCanvas as unknown as LGraphCanvas,
      dataTransfer.items,
      mockNode as unknown as LGraphNode
    )

    expect(mockNode.pasteFile).toHaveBeenCalledWith(file1)
    expect(mockNode.pasteFiles).toHaveBeenCalledWith([file1, file2])
  })

  it('should do nothing when no image files present', async () => {
    const mockNode = createMockNode()
    const dataTransfer = createDataTransfer()

    await pasteImageNode(
      mockCanvas as unknown as LGraphCanvas,
      dataTransfer.items,
      mockNode as unknown as LGraphNode
    )

    expect(mockNode.pasteFile).not.toHaveBeenCalled()
    expect(mockNode.pasteFiles).not.toHaveBeenCalled()
  })

  it('should filter non-image items', async () => {
    const mockNode = createMockNode()
    const imageFile = createImageFile()
    const textFile = new File([''], 'test.txt', { type: 'text/plain' })
    const dataTransfer = createDataTransfer([textFile, imageFile])

    await pasteImageNode(
      mockCanvas as unknown as LGraphCanvas,
      dataTransfer.items,
      mockNode as unknown as LGraphNode
    )

    expect(mockNode.pasteFile).toHaveBeenCalledWith(imageFile)
    expect(mockNode.pasteFiles).toHaveBeenCalledWith([imageFile])
  })
})

describe('pasteImageNodes', () => {
  beforeEach(() => {
    vi.clearAllMocks()
  })

  it('should create multiple nodes for multiple files', async () => {
    const mockNode1 = createMockNode()
    const mockNode2 = createMockNode()
    vi.mocked(createNode)
      .mockResolvedValueOnce(mockNode1 as unknown as LGraphNode)
      .mockResolvedValueOnce(mockNode2 as unknown as LGraphNode)

    const file1 = createImageFile('test1.png')
    const file2 = createImageFile('test2.jpg', 'image/jpeg')

    const result = await pasteImageNodes(
      mockCanvas as unknown as LGraphCanvas,
      [file1, file2]
    )

    expect(createNode).toHaveBeenCalledTimes(2)
    expect(createNode).toHaveBeenNthCalledWith(1, mockCanvas, 'LoadImage')
    expect(createNode).toHaveBeenNthCalledWith(2, mockCanvas, 'LoadImage')
    expect(mockNode1.pasteFile).toHaveBeenCalledWith(file1)
    expect(mockNode2.pasteFile).toHaveBeenCalledWith(file2)
    expect(result).toEqual([mockNode1, mockNode2])
  })

  it('should handle empty file list', async () => {
    const result = await pasteImageNodes(
      mockCanvas as unknown as LGraphCanvas,
      []
    )

    expect(createNode).not.toHaveBeenCalled()
    expect(result).toEqual([])
  })
})

describe('usePaste', () => {
  beforeEach(() => {
    vi.clearAllMocks()
    mockCanvas.current_node = null
    mockWorkspaceStore.shiftDown = false
    vi.mocked(mockCanvas.graph!.add).mockImplementation(
      (node: LGraphNode | LGraphGroup | null) => node as LGraphNode
    )
  })

  it('should handle image paste', async () => {
    const mockNode = createMockNode()
    vi.mocked(createNode).mockResolvedValue(mockNode as unknown as LGraphNode)

    usePaste()

    const file = createImageFile()
    const dataTransfer = createDataTransfer([file])
    const event = new ClipboardEvent('paste', { clipboardData: dataTransfer })
    document.dispatchEvent(event)

    await vi.waitFor(() => {
      expect(createNode).toHaveBeenCalledWith(mockCanvas, 'LoadImage')
      expect(mockNode.pasteFile).toHaveBeenCalledWith(file)
    })
  })

  it('should handle audio paste', async () => {
    const mockNode = createMockNode()
    vi.mocked(LiteGraph.createNode).mockReturnValue(
      mockNode as unknown as LGraphNode
    )

    usePaste()

    const file = createAudioFile()
    const dataTransfer = createDataTransfer([file])
    const event = new ClipboardEvent('paste', { clipboardData: dataTransfer })
    document.dispatchEvent(event)

    await vi.waitFor(() => {
      expect(LiteGraph.createNode).toHaveBeenCalledWith('LoadAudio')
      expect(mockNode.pasteFile).toHaveBeenCalledWith(file)
    })
  })

  it('should handle workflow JSON paste', async () => {
    const workflow = { version: '1.0', nodes: [], extra: {} }

    usePaste()

    const dataTransfer = new DataTransfer()
    dataTransfer.setData('text/plain', JSON.stringify(workflow))

    const event = new ClipboardEvent('paste', { clipboardData: dataTransfer })
    document.dispatchEvent(event)

    await vi.waitFor(() => {
      expect(app.loadGraphData).toHaveBeenCalledWith(workflow)
    })
  })

  it('should ignore paste when shift is down', () => {
    mockWorkspaceStore.shiftDown = true

    usePaste()

    const file = createImageFile()
    const dataTransfer = createDataTransfer([file])
    const event = new ClipboardEvent('paste', { clipboardData: dataTransfer })
    document.dispatchEvent(event)

    expect(LiteGraph.createNode).not.toHaveBeenCalled()
  })

  it('should use existing image node when selected', () => {
    const mockNode = {
      is_selected: true,
      pasteFile: vi.fn(),
      pasteFiles: vi.fn()
    } as unknown as Partial<LGraphNode> as LGraphNode
    mockCanvas.current_node = mockNode
    vi.mocked(isImageNode).mockReturnValue(true)

    usePaste()

    const file = createImageFile()
    const dataTransfer = createDataTransfer([file])
    const event = new ClipboardEvent('paste', { clipboardData: dataTransfer })
    document.dispatchEvent(event)

    expect(mockNode.pasteFile).toHaveBeenCalledWith(file)
  })

  it('should call canvas pasteFromClipboard for non-workflow text', () => {
    usePaste()

    const dataTransfer = new DataTransfer()
    dataTransfer.setData('text/plain', 'just some text')

    const event = new ClipboardEvent('paste', { clipboardData: dataTransfer })
    document.dispatchEvent(event)

    expect(mockCanvas.pasteFromClipboard).toHaveBeenCalled()
  })

  it('should handle clipboard items with metadata', async () => {
    const data = { test: 'data' }
    const encoded = btoa(JSON.stringify(data))
    const html = `<div data-metadata="${encoded}"></div>`

    usePaste()

    const dataTransfer = new DataTransfer()
    dataTransfer.setData('text/html', html)

    const event = new ClipboardEvent('paste', { clipboardData: dataTransfer })
    document.dispatchEvent(event)

    await vi.waitFor(() => {
      expect(mockCanvas._deserializeItems).toHaveBeenCalledWith(
        data,
        expect.any(Object)
      )
    })
  })
})

describe('cloneDataTransfer', () => {
  it('should clone string data', () => {
    const original = new DataTransfer()
    original.setData('text/plain', 'test text')
    original.setData('text/html', '<p>test html</p>')

    const cloned = cloneDataTransfer(original)

    expect(cloned.getData('text/plain')).toBe('test text')
    expect(cloned.getData('text/html')).toBe('<p>test html</p>')
  })

  it('should clone files', () => {
    const file1 = createImageFile('test1.png')
    const file2 = createImageFile('test2.jpg', 'image/jpeg')
    const original = createDataTransfer([file1, file2])

    const cloned = cloneDataTransfer(original)

    // Files are added from both .files and .items, causing duplicates
    expect(cloned.files.length).toBeGreaterThanOrEqual(2)
    expect(Array.from(cloned.files)).toContain(file1)
    expect(Array.from(cloned.files)).toContain(file2)
  })

  it('should preserve dropEffect and effectAllowed', () => {
    const original = new DataTransfer()
    original.dropEffect = 'copy'
    original.effectAllowed = 'copyMove'

    const cloned = cloneDataTransfer(original)

    expect(cloned.dropEffect).toBe('copy')
    expect(cloned.effectAllowed).toBe('copyMove')
  })

  it('should handle empty DataTransfer', () => {
    const original = new DataTransfer()

    const cloned = cloneDataTransfer(original)

    expect(cloned.types.length).toBe(0)
    expect(cloned.files.length).toBe(0)
  })

  it('should clone both string data and files', () => {
    const file = createImageFile()
    const original = createDataTransfer([file])
    original.setData('text/plain', 'test')

    const cloned = cloneDataTransfer(original)

    expect(cloned.getData('text/plain')).toBe('test')
    // Files are added from both .files and .items
    expect(cloned.files.length).toBeGreaterThanOrEqual(1)
    expect(Array.from(cloned.files)).toContain(file)
  })
})