ComfyUI_frontend/src/services/jobOutputCache.test.ts

import { beforeEach, describe, expect, it, vi } from 'vitest'

import type {
  JobDetail,
  JobListItem
} from '@/platform/remote/comfyui/jobs/jobTypes'
import { ResultItemImpl, TaskItemImpl } from '@/stores/queueStore'

vi.mock('@/platform/remote/comfyui/jobs/fetchJobs', () => ({
  fetchJobDetail: vi.fn(),
  extractWorkflow: vi.fn()
}))

function createResultItem(url: string, supportsPreview = true): ResultItemImpl {
  const item = new ResultItemImpl({
    filename: url,
    subfolder: '',
    type: 'output',
    nodeId: 'node-1',
    mediaType: supportsPreview ? 'images' : 'unknown'
  })
  Object.defineProperty(item, 'url', { get: () => url })
  Object.defineProperty(item, 'supportsPreview', { get: () => supportsPreview })
  return item
}

function createMockJob(id: string, outputsCount = 1): JobListItem {
  return {
    id,
    status: 'completed',
    create_time: Date.now(),
    preview_output: null,
    outputs_count: outputsCount,
    priority: 0
  }
}

function createTask(
  preview?: ResultItemImpl,
  allOutputs?: ResultItemImpl[],
  outputsCount = 1
): TaskItemImpl {
  const job = createMockJob(
    `task-${Math.random().toString(36).slice(2)}`,
    outputsCount
  )
  const flatOutputs = allOutputs ?? (preview ? [preview] : [])
  return new TaskItemImpl(job, {}, flatOutputs)
}

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

  describe('findActiveIndex', () => {
    it('returns index of matching URL', async () => {
      const { findActiveIndex } = await import('@/services/jobOutputCache')
      const items = [
        createResultItem('a'),
        createResultItem('b'),
        createResultItem('c')
      ]

      expect(findActiveIndex(items, 'b')).toBe(1)
    })

    it('returns 0 when URL not found', async () => {
      const { findActiveIndex } = await import('@/services/jobOutputCache')
      const items = [createResultItem('a'), createResultItem('b')]

      expect(findActiveIndex(items, 'missing')).toBe(0)
    })

    it('returns 0 when URL is undefined', async () => {
      const { findActiveIndex } = await import('@/services/jobOutputCache')
      const items = [createResultItem('a'), createResultItem('b')]

      expect(findActiveIndex(items, undefined)).toBe(0)
    })
  })

  describe('getOutputsForTask', () => {
    it('returns previewable outputs directly when no lazy load needed', async () => {
      const { getOutputsForTask } = await import('@/services/jobOutputCache')
      const outputs = [createResultItem('p-1'), createResultItem('p-2')]
      const task = createTask(undefined, outputs, 1)

      const result = await getOutputsForTask(task)

      expect(result).toEqual(outputs)
    })

    it('lazy loads when outputsCount > 1', async () => {
      const { getOutputsForTask } = await import('@/services/jobOutputCache')
      const previewOutput = createResultItem('preview')
      const fullOutputs = [
        createResultItem('full-1'),
        createResultItem('full-2')
      ]

      const job = createMockJob('task-1', 3)
      const task = new TaskItemImpl(job, {}, [previewOutput])
      const loadedTask = new TaskItemImpl(job, {}, fullOutputs)
      task.loadFullOutputs = vi.fn().mockResolvedValue(loadedTask)

      const result = await getOutputsForTask(task)

      expect(result).toEqual(fullOutputs)
      expect(task.loadFullOutputs).toHaveBeenCalled()
    })

    it('caches loaded tasks', async () => {
      const { getOutputsForTask } = await import('@/services/jobOutputCache')
      const fullOutputs = [createResultItem('full-1')]

      const job = createMockJob('task-1', 3)
      const task = new TaskItemImpl(job, {}, [createResultItem('preview')])
      const loadedTask = new TaskItemImpl(job, {}, fullOutputs)
      task.loadFullOutputs = vi.fn().mockResolvedValue(loadedTask)

      // First call should load
      await getOutputsForTask(task)
      expect(task.loadFullOutputs).toHaveBeenCalledTimes(1)

      // Second call should use cache
      await getOutputsForTask(task)
      expect(task.loadFullOutputs).toHaveBeenCalledTimes(1)
    })

    it('falls back to preview outputs on load error', async () => {
      const { getOutputsForTask } = await import('@/services/jobOutputCache')
      const previewOutput = createResultItem('preview')

      const job = createMockJob('task-1', 3)
      const task = new TaskItemImpl(job, {}, [previewOutput])
      task.loadFullOutputs = vi
        .fn()
        .mockRejectedValue(new Error('Network error'))

      const result = await getOutputsForTask(task)

      expect(result).toEqual([previewOutput])
    })

    it('returns null when request is superseded', async () => {
      const { getOutputsForTask } = await import('@/services/jobOutputCache')
      const job1 = createMockJob('task-1', 3)
      const job2 = createMockJob('task-2', 3)

      const task1 = new TaskItemImpl(job1, {}, [createResultItem('preview-1')])
      const task2 = new TaskItemImpl(job2, {}, [createResultItem('preview-2')])

      const loadedTask1 = new TaskItemImpl(job1, {}, [
        createResultItem('full-1')
      ])
      const loadedTask2 = new TaskItemImpl(job2, {}, [
        createResultItem('full-2')
      ])

      // Task1 loads slowly, task2 loads quickly
      task1.loadFullOutputs = vi.fn().mockImplementation(
        () =>
          new Promise((resolve) => {
            setTimeout(() => resolve(loadedTask1), 50)
          })
      )
      task2.loadFullOutputs = vi.fn().mockResolvedValue(loadedTask2)

      // Start task1, then immediately start task2
      const promise1 = getOutputsForTask(task1)
      const promise2 = getOutputsForTask(task2)

      const [result1, result2] = await Promise.all([promise1, promise2])

      // Task2 should succeed, task1 should return null (superseded)
      expect(result1).toBeNull()
      expect(result2).toEqual([createResultItem('full-2')])
    })
  })

  describe('getJobDetail', () => {
    it('fetches and caches job detail', async () => {
      const { getJobDetail } = await import('@/services/jobOutputCache')
      const { fetchJobDetail } =
        await import('@/platform/remote/comfyui/jobs/fetchJobs')

      const mockDetail: JobDetail = {
        id: 'job-1',
        status: 'completed',
        create_time: Date.now(),
        priority: 0,
        outputs: {}
      }
      vi.mocked(fetchJobDetail).mockResolvedValue(mockDetail)

      const result = await getJobDetail('job-1')

      expect(result).toEqual(mockDetail)
      expect(fetchJobDetail).toHaveBeenCalledWith(expect.any(Function), 'job-1')
    })

    it('returns cached job detail on subsequent calls', async () => {
      const { getJobDetail } = await import('@/services/jobOutputCache')
      const { fetchJobDetail } =
        await import('@/platform/remote/comfyui/jobs/fetchJobs')

      const mockDetail: JobDetail = {
        id: 'job-2',
        status: 'completed',
        create_time: Date.now(),
        priority: 0,
        outputs: {}
      }
      vi.mocked(fetchJobDetail).mockResolvedValue(mockDetail)

      // First call
      await getJobDetail('job-2')
      expect(fetchJobDetail).toHaveBeenCalledTimes(1)

      // Second call should use cache
      const result = await getJobDetail('job-2')
      expect(result).toEqual(mockDetail)
      expect(fetchJobDetail).toHaveBeenCalledTimes(1)
    })

    it('returns undefined on fetch error', async () => {
      const { getJobDetail } = await import('@/services/jobOutputCache')
      const { fetchJobDetail } =
        await import('@/platform/remote/comfyui/jobs/fetchJobs')

      vi.mocked(fetchJobDetail).mockRejectedValue(new Error('Network error'))

      const result = await getJobDetail('job-error')

      expect(result).toBeUndefined()
    })
  })

  describe('getJobWorkflow', () => {
    it('fetches job detail and extracts workflow', async () => {
      const { getJobWorkflow } = await import('@/services/jobOutputCache')
      const { fetchJobDetail, extractWorkflow } =
        await import('@/platform/remote/comfyui/jobs/fetchJobs')

      const mockDetail: JobDetail = {
        id: 'job-wf',
        status: 'completed',
        create_time: Date.now(),
        priority: 0,
        outputs: {}
      }
      const mockWorkflow = { version: 1 }

      vi.mocked(fetchJobDetail).mockResolvedValue(mockDetail)
      vi.mocked(extractWorkflow).mockResolvedValue(mockWorkflow as any)

      const result = await getJobWorkflow('job-wf')

      expect(result).toEqual(mockWorkflow)
      expect(extractWorkflow).toHaveBeenCalledWith(mockDetail)
    })

    it('returns undefined when job detail not found', async () => {
      const { getJobWorkflow } = await import('@/services/jobOutputCache')
      const { fetchJobDetail, extractWorkflow } =
        await import('@/platform/remote/comfyui/jobs/fetchJobs')

      vi.mocked(fetchJobDetail).mockResolvedValue(undefined)
      vi.mocked(extractWorkflow).mockResolvedValue(undefined)

      const result = await getJobWorkflow('missing')

      expect(result).toBeUndefined()
    })
  })
})