ComfyUI_frontend/src/stores/resultItemParsing.test.ts

import { fromPartial } from '@total-typescript/shoehorn'
import { describe, expect, it } from 'vitest'

import type { NodeExecutionOutput } from '@/schemas/apiSchema'
import { parseNodeOutput, parseTaskOutput } from '@/stores/resultItemParsing'

function makeOutput(
  overrides: Partial<NodeExecutionOutput> = {}
): NodeExecutionOutput {
  return { ...overrides }
}

describe(parseNodeOutput, () => {
  it('returns empty array for output with no known media types', () => {
    const result = parseNodeOutput('1', makeOutput({ text: 'hello' }))
    expect(result).toEqual([])
  })

  it('flattens images into ResultItemImpl instances', () => {
    const output = makeOutput({
      images: [
        { filename: 'a.png', subfolder: '', type: 'output' },
        { filename: 'b.png', subfolder: 'sub', type: 'output' }
      ]
    })

    const result = parseNodeOutput('42', output)

    expect(result).toHaveLength(2)
    expect(result[0].filename).toBe('a.png')
    expect(result[0].nodeId).toBe('42')
    expect(result[0].mediaType).toBe('images')
    expect(result[1].filename).toBe('b.png')
    expect(result[1].subfolder).toBe('sub')
  })

  it('flattens audio outputs', () => {
    const output = makeOutput({
      audio: [{ filename: 'sound.wav', subfolder: '', type: 'output' }]
    })

    const result = parseNodeOutput(7, output)

    expect(result).toHaveLength(1)
    expect(result[0].mediaType).toBe('audio')
    expect(result[0].nodeId).toBe(7)
  })

  it('flattens multiple media types in a single output', () => {
    const output = makeOutput({
      images: [{ filename: 'img.png', subfolder: '', type: 'output' }],
      video: [{ filename: 'vid.mp4', subfolder: '', type: 'output' }]
    })

    const result = parseNodeOutput('1', output)

    expect(result).toHaveLength(2)
    const types = result.map((r) => r.mediaType)
    expect(types).toContain('images')
    expect(types).toContain('video')
  })

  it('handles gifs and 3d output types', () => {
    const output = makeOutput({
      gifs: [
        { filename: 'anim.gif', subfolder: '', type: 'output' }
      ] as NodeExecutionOutput['gifs'],
      '3d': [
        { filename: 'model.glb', subfolder: '', type: 'output' }
      ] as NodeExecutionOutput['3d']
    })

    const result = parseNodeOutput('5', output)

    expect(result).toHaveLength(2)
    const types = result.map((r) => r.mediaType)
    expect(types).toContain('gifs')
    expect(types).toContain('3d')
  })

  it('ignores empty arrays', () => {
    const output = makeOutput({ images: [], audio: [] })
    const result = parseNodeOutput('1', output)
    expect(result).toEqual([])
  })

  it('excludes animated key', () => {
    const output = makeOutput({
      images: [{ filename: 'img.png', subfolder: '', type: 'output' }],
      animated: [true]
    })

    const result = parseNodeOutput('1', output)

    expect(result).toHaveLength(1)
    expect(result[0].mediaType).toBe('images')
  })

  it('excludes text key', () => {
    const output = makeOutput({
      images: [{ filename: 'img.png', subfolder: '', type: 'output' }],
      text: 'some text output'
    })

    const result = parseNodeOutput('1', output)

    expect(result).toHaveLength(1)
    expect(result[0].mediaType).toBe('images')
  })

  it('excludes non-ResultItem array items', () => {
    const output = fromPartial<NodeExecutionOutput>({
      images: [{ filename: 'img.png', subfolder: '', type: 'output' }],
      custom_data: [{ randomKey: 123 }]
    })

    const result = parseNodeOutput('1', output)

    expect(result).toHaveLength(1)
    expect(result[0].mediaType).toBe('images')
  })

  it('accepts items with filename but no subfolder', () => {
    const output = fromPartial<NodeExecutionOutput>({
      images: [
        { filename: 'valid.png', subfolder: '', type: 'output' },
        { filename: 'no-subfolder.png' }
      ]
    })

    const result = parseNodeOutput('1', output)

    expect(result).toHaveLength(2)
    expect(result[0].filename).toBe('valid.png')
    expect(result[1].filename).toBe('no-subfolder.png')
    expect(result[1].subfolder).toBe('')
  })

  it('excludes items missing filename', () => {
    const output = fromPartial<NodeExecutionOutput>({
      images: [
        { filename: 'valid.png', subfolder: '', type: 'output' },
        { subfolder: '', type: 'output' }
      ]
    })

    const result = parseNodeOutput('1', output)

    expect(result).toHaveLength(1)
    expect(result[0].filename).toBe('valid.png')
  })
})

describe(parseTaskOutput, () => {
  it('flattens across multiple nodes', () => {
    const taskOutput: Record<string, NodeExecutionOutput> = {
      '1': makeOutput({
        images: [{ filename: 'a.png', subfolder: '', type: 'output' }]
      }),
      '2': makeOutput({
        audio: [{ filename: 'b.wav', subfolder: '', type: 'output' }]
      })
    }

    const result = parseTaskOutput(taskOutput)

    expect(result).toHaveLength(2)
    expect(result[0].nodeId).toBe('1')
    expect(result[0].filename).toBe('a.png')
    expect(result[1].nodeId).toBe('2')
    expect(result[1].filename).toBe('b.wav')
  })
})