ComfyUI_frontend/browser_tests/fixtures/helpers/AssetsHelper.ts

import type { Page, Route } from '@playwright/test'
import type {
  CreateAssetExportData,
  CreateAssetExportResponse,
  JobsListResponse,
  ListAssetsResponse
} from '@comfyorg/ingest-types'

import type {
  JobDetail,
  RawJobListItem
} from '@/platform/remote/comfyui/jobs/jobTypes'

const jobsListRoutePattern = '**/api/jobs?*'
const assetsListRoutePattern = /\/api\/assets(?:\?.*)?$/
const assetExportRoutePattern = '**/api/assets/export'
const inputFilesRoutePattern = /\/internal\/files\/input(?:\?.*)?$/
const historyRoutePattern = /\/api\/history$/

/**
 * Media kinds supported by the assets sidebar filter UI. The string values
 * match what the backend stores on `preview_output.mediaType` (`images` is
 * intentionally plural to match existing API conventions; the others are
 * singular as emitted by `useMediaAssetGalleryStore`).
 *
 * The sidebar filter ultimately matches on the filename extension, so the
 * fixture also picks an extension-appropriate filename for each kind.
 */
export type MediaKindFixture = 'images' | 'video' | 'audio' | '3D'

const DEFAULT_EXTENSION: Record<MediaKindFixture, string> = {
  images: 'png',
  video: 'mp4',
  audio: 'wav',
  '3D': 'glb'
}

/** Factory to create a mock completed job with preview output. */
export function createMockJob(
  overrides: Partial<RawJobListItem> & {
    id: string
    /**
     * Optional shorthand to set both `preview_output.mediaType` and an
     * extension-appropriate filename. Ignored when `preview_output` is also
     * supplied via `overrides`.
     */
    mediaKind?: MediaKindFixture
  }
): RawJobListItem {
  const { mediaKind, ...rest } = overrides
  const now = Date.now()
  const extension = mediaKind ? DEFAULT_EXTENSION[mediaKind] : 'png'
  const mediaType = mediaKind ?? 'images'

  return {
    status: 'completed',
    create_time: now,
    execution_start_time: now,
    execution_end_time: now + 5000,
    preview_output: {
      filename: `output_${rest.id}.${extension}`,
      subfolder: '',
      type: 'output',
      nodeId: '1',
      mediaType
    },
    outputs_count: 1,
    priority: 0,
    ...rest
  }
}

/** Create multiple mock jobs with sequential IDs and staggered timestamps. */
export function createMockJobs(
  count: number,
  baseOverrides?: Partial<RawJobListItem>
): RawJobListItem[] {
  const now = Date.now()
  return Array.from({ length: count }, (_, i) =>
    createMockJob({
      id: `job-${String(i + 1).padStart(3, '0')}`,
      create_time: now - i * 60_000,
      execution_start_time: now - i * 60_000,
      execution_end_time: now - i * 60_000 + 5000 + i * 1000,
      preview_output: {
        filename: `image_${String(i + 1).padStart(3, '0')}.png`,
        subfolder: '',
        type: 'output',
        nodeId: '1',
        mediaType: 'images'
      },
      ...baseOverrides
    })
  )
}

/**
 * Create one job per requested media kind, in the order supplied. Jobs share
 * a stable timestamp ordering (newer first) so callers can rely on the result
 * order when mediaType filters are inactive.
 */
export function createMixedMediaJobs(
  kinds: MediaKindFixture[]
): RawJobListItem[] {
  const now = Date.now()
  return kinds.map((kind, i) =>
    createMockJob({
      id: `${kind}-${String(i + 1).padStart(3, '0')}`,
      mediaKind: kind,
      create_time: now - i * 60_000,
      execution_start_time: now - i * 60_000,
      execution_end_time: now - i * 60_000 + 5000
    })
  )
}

/**
 * Create jobs with explicit `(create_time, execution duration)` pairs so that
 * sort assertions for newest/oldest and longest/fastest are unambiguous.
 *
 * Each spec entry yields a job whose `execution_end_time - execution_start_time`
 * equals `durationMs`. The first spec becomes id `job-001`, etc.
 */
export function createJobsWithExecutionTimes(
  specs: ReadonlyArray<{ createTime: number; durationMs: number }>
): RawJobListItem[] {
  return specs.map((spec, i) =>
    createMockJob({
      id: `job-${String(i + 1).padStart(3, '0')}`,
      create_time: spec.createTime,
      execution_start_time: spec.createTime,
      execution_end_time: spec.createTime + spec.durationMs
    })
  )
}

/** Create mock imported file names with various media types. */
export function createMockImportedFiles(count: number): string[] {
  const extensions = ['png', 'jpg', 'mp4', 'wav', 'glb', 'txt']
  return Array.from(
    { length: count },
    (_, i) =>
      `imported_${String(i + 1).padStart(3, '0')}.${extensions[i % extensions.length]}`
  )
}

function parseLimit(url: URL, total: number): number {
  const value = Number(url.searchParams.get('limit'))
  if (!Number.isInteger(value) || value <= 0) {
    return total
  }
  return value
}

function parseOffset(url: URL): number {
  const value = Number(url.searchParams.get('offset'))
  if (!Number.isInteger(value) || value < 0) {
    return 0
  }
  return value
}

function getExecutionDuration(job: RawJobListItem): number {
  const start = job.execution_start_time ?? 0
  const end = job.execution_end_time ?? 0
  return end - start
}

export class AssetsHelper {
  private jobsRouteHandler: ((route: Route) => Promise<void>) | null = null
  private cloudAssetsRouteHandler: ((route: Route) => Promise<void>) | null =
    null
  private assetExportRouteHandler: ((route: Route) => Promise<void>) | null =
    null
  private inputFilesRouteHandler: ((route: Route) => Promise<void>) | null =
    null
  private deleteHistoryRouteHandler: ((route: Route) => Promise<void>) | null =
    null
  private generatedJobs: RawJobListItem[] = []
  private cloudAssetsResponse: ListAssetsResponse | null = null
  private assetExportRequests: CreateAssetExportData['body'][] = []
  private assetExportResponse: CreateAssetExportResponse | null = null
  private importedFiles: string[] = []
  private readonly jobDetailRouteHandlers = new Map<
    string,
    (route: Route) => Promise<void>
  >()

  constructor(private readonly page: Page) {}

  async mockOutputHistory(jobs: RawJobListItem[]): Promise<void> {
    this.generatedJobs = [...jobs]

    if (this.jobsRouteHandler) {
      return
    }

    this.jobsRouteHandler = async (route: Route) => {
      const url = new URL(route.request().url())
      const statuses = url.searchParams
        .get('status')
        ?.split(',')
        .map((status) => status.trim())
        .filter(Boolean)
      const workflowId = url.searchParams.get('workflow_id')
      const sortBy = url.searchParams.get('sort_by')
      const sortOrder = url.searchParams.get('sort_order') === 'asc' ? 1 : -1

      let filteredJobs = [...this.generatedJobs]

      if (statuses?.length) {
        filteredJobs = filteredJobs.filter((job) =>
          statuses.includes(job.status)
        )
      }

      if (workflowId) {
        filteredJobs = filteredJobs.filter(
          (job) => job.workflow_id === workflowId
        )
      }

      filteredJobs.sort((left, right) => {
        const leftValue =
          sortBy === 'execution_duration'
            ? getExecutionDuration(left)
            : left.create_time
        const rightValue =
          sortBy === 'execution_duration'
            ? getExecutionDuration(right)
            : right.create_time

        return (leftValue - rightValue) * sortOrder
      })

      const offset = parseOffset(url)
      const total = filteredJobs.length
      const limit = parseLimit(url, total)
      const visibleJobs = filteredJobs.slice(offset, offset + limit)

      const response = {
        jobs: visibleJobs,
        pagination: {
          offset,
          limit,
          total,
          has_more: offset + visibleJobs.length < total
        }
      } satisfies {
        jobs: unknown[]
        pagination: JobsListResponse['pagination']
      }

      await route.fulfill({
        status: 200,
        contentType: 'application/json',
        body: JSON.stringify(response)
      })
    }

    await this.page.route(jobsListRoutePattern, this.jobsRouteHandler)
  }

  async mockCloudAssets(response: ListAssetsResponse): Promise<void> {
    this.cloudAssetsResponse = response

    if (this.cloudAssetsRouteHandler) {
      return
    }

    this.cloudAssetsRouteHandler = async (route: Route) => {
      await route.fulfill({
        status: 200,
        contentType: 'application/json',
        body: JSON.stringify(this.cloudAssetsResponse)
      })
    }

    await this.page.route(assetsListRoutePattern, this.cloudAssetsRouteHandler)
  }

  async mockEmptyCloudAssets(): Promise<void> {
    await this.mockCloudAssets({
      assets: [],
      total: 0,
      has_more: false
    })
  }

  async captureAssetExportRequests(
    response: CreateAssetExportResponse = {
      task_id: 'asset-export-task',
      status: 'created'
    }
  ): Promise<CreateAssetExportData['body'][]> {
    this.assetExportRequests = []
    this.assetExportResponse = response

    if (this.assetExportRouteHandler) {
      return this.assetExportRequests
    }

    this.assetExportRouteHandler = async (route: Route) => {
      this.assetExportRequests.push(
        route.request().postDataJSON() as CreateAssetExportData['body']
      )

      await route.fulfill({
        status: 202,
        contentType: 'application/json',
        body: JSON.stringify(this.assetExportResponse)
      })
    }

    await this.page.route(assetExportRoutePattern, this.assetExportRouteHandler)

    return this.assetExportRequests
  }

  async mockJobDetail(jobId: string, detail: JobDetail): Promise<void> {
    const pattern = `**/api/jobs/${encodeURIComponent(jobId)}`
    const existingHandler = this.jobDetailRouteHandlers.get(pattern)

    if (existingHandler) {
      await this.page.unroute(pattern, existingHandler)
    }

    const handler = async (route: Route) => {
      await route.fulfill({
        status: 200,
        contentType: 'application/json',
        body: JSON.stringify(detail)
      })
    }

    this.jobDetailRouteHandlers.set(pattern, handler)
    await this.page.route(pattern, handler)
  }

  async mockInputFiles(files: string[]): Promise<void> {
    this.importedFiles = [...files]

    if (this.inputFilesRouteHandler) {
      return
    }

    this.inputFilesRouteHandler = async (route: Route) => {
      await route.fulfill({
        status: 200,
        contentType: 'application/json',
        body: JSON.stringify(this.importedFiles)
      })
    }

    await this.page.route(inputFilesRoutePattern, this.inputFilesRouteHandler)
  }

  /**
   * Mock the POST /api/history endpoint used for deleting history items.
   * On receiving a `{ delete: [id] }` payload, removes matching jobs from
   * the in-memory mock state so subsequent /api/jobs fetches reflect the
   * deletion.
   */
  async mockDeleteHistory(): Promise<void> {
    if (this.deleteHistoryRouteHandler) return

    this.deleteHistoryRouteHandler = async (route: Route) => {
      const request = route.request()
      if (request.method() !== 'POST') {
        await route.continue()
        return
      }

      const body = request.postDataJSON() as { delete?: string[] }
      if (body.delete) {
        const idsToRemove = new Set(body.delete)
        this.generatedJobs = this.generatedJobs.filter(
          (job) => !idsToRemove.has(job.id)
        )
      }

      await route.fulfill({ status: 200, body: '{}' })
    }

    await this.page.route(historyRoutePattern, this.deleteHistoryRouteHandler)
  }

  async mockEmptyState(): Promise<void> {
    await this.mockOutputHistory([])
    await this.mockInputFiles([])
  }

  async clearMocks(): Promise<void> {
    this.generatedJobs = []
    this.cloudAssetsResponse = null
    this.assetExportRequests = []
    this.assetExportResponse = null
    this.importedFiles = []

    if (this.jobsRouteHandler) {
      await this.page.unroute(jobsListRoutePattern, this.jobsRouteHandler)
      this.jobsRouteHandler = null
    }

    if (this.cloudAssetsRouteHandler) {
      await this.page.unroute(
        assetsListRoutePattern,
        this.cloudAssetsRouteHandler
      )
      this.cloudAssetsRouteHandler = null
    }

    if (this.assetExportRouteHandler) {
      await this.page.unroute(
        assetExportRoutePattern,
        this.assetExportRouteHandler
      )
      this.assetExportRouteHandler = null
    }

    if (this.inputFilesRouteHandler) {
      await this.page.unroute(
        inputFilesRoutePattern,
        this.inputFilesRouteHandler
      )
      this.inputFilesRouteHandler = null
    }

    if (this.deleteHistoryRouteHandler) {
      await this.page.unroute(
        historyRoutePattern,
        this.deleteHistoryRouteHandler
      )
      this.deleteHistoryRouteHandler = null
    }

    for (const [pattern, handler] of this.jobDetailRouteHandlers) {
      await this.page.unroute(pattern, handler)
    }
    this.jobDetailRouteHandlers.clear()
  }
}