ComfyUI_frontend/browser_tests/fixtures/helpers/ExecutionHelper.ts

import type { Page } from '@playwright/test'

import type { RawJobListItem } from '../../../src/platform/remote/comfyui/jobs/jobTypes'
import type { MockWebSocket } from '../ws'

/**
 * Helper for simulating prompt execution in e2e tests.
 */
export class ExecutionHelper {
  private jobCounter = 0

  constructor(
    private readonly page: Page,
    private readonly mock: MockWebSocket
  ) {}

  private get ws() {
    return this.mock.ws
  }

  /**
   * Intercept POST /api/prompt, execute Comfy.QueuePrompt, and return
   * the synthetic job ID.
   *
   * The app receives a valid PromptResponse so storeJob() fires
   * and registers the job against the active workflow path.
   */
  async run(): Promise<string> {
    const jobId = `test-job-${++this.jobCounter}`

    let fulfilled!: () => void
    const prompted = new Promise<void>((r) => {
      fulfilled = r
    })

    await this.page.route(
      '**/api/prompt',
      async (route) => {
        await route.fulfill({
          status: 200,
          contentType: 'application/json',
          body: JSON.stringify({
            prompt_id: jobId,
            node_errors: {}
          })
        })
        fulfilled()
      },
      { times: 1 }
    )

    await this.page.evaluate(() => {
      window.app!.extensionManager.command.execute('Comfy.QueuePrompt')
    })
    await prompted

    // Prevent real server events from interfering with test-controlled execution
    this.mock.stopServerForwarding()

    return jobId
  }

  /**
   * Send a binary `b_preview_with_metadata` WS message (type 4).
   * Encodes the metadata and a tiny 1x1 PNG so the app creates a blob URL.
   */
  latentPreview(jobId: string, nodeId: string): void {
    const metadata = JSON.stringify({
      node_id: nodeId,
      display_node_id: nodeId,
      parent_node_id: nodeId,
      real_node_id: nodeId,
      prompt_id: jobId,
      image_type: 'image/png'
    })
    const metadataBytes = new TextEncoder().encode(metadata)

    // 1x1 red PNG
    const png = Buffer.from(
      'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwADhQGAWjR9awAAAABJRU5ErkJggg==',
      'base64'
    )

    // Binary format: [type:uint32][metadataLength:uint32][metadata][imageData]
    const buf = new ArrayBuffer(8 + metadataBytes.length + png.length)
    const view = new DataView(buf)
    view.setUint32(0, 4) // type 4 = PREVIEW_IMAGE_WITH_METADATA
    view.setUint32(4, metadataBytes.length)
    new Uint8Array(buf, 8, metadataBytes.length).set(metadataBytes)
    new Uint8Array(buf, 8 + metadataBytes.length).set(png)

    this.ws.send(Buffer.from(buf))
  }

  /** Send `execution_start` WS event. */
  executionStart(jobId: string): void {
    this.ws.send(
      JSON.stringify({
        type: 'execution_start',
        data: { prompt_id: jobId, timestamp: Date.now() }
      })
    )
  }

  /** Send `executed` WS event with node output. */
  executed(
    jobId: string,
    nodeId: string,
    output: Record<string, unknown>
  ): void {
    this.ws.send(
      JSON.stringify({
        type: 'executed',
        data: {
          prompt_id: jobId,
          node: nodeId,
          display_node: nodeId,
          output
        }
      })
    )
  }

  /** Send `execution_success` WS event. */
  executionSuccess(jobId: string): void {
    this.ws.send(
      JSON.stringify({
        type: 'execution_success',
        data: { prompt_id: jobId, timestamp: Date.now() }
      })
    )
  }

  /** Send `execution_error` WS event. */
  executionError(jobId: string, nodeId: string, message: string): void {
    this.ws.send(
      JSON.stringify({
        type: 'execution_error',
        data: {
          prompt_id: jobId,
          timestamp: Date.now(),
          node_id: nodeId,
          node_type: 'Unknown',
          exception_message: message,
          exception_type: 'RuntimeError',
          traceback: []
        }
      })
    )
  }

  /** Send `progress` WS event. */
  progress(jobId: string, nodeId: string, value: number, max: number): void {
    this.ws.send(
      JSON.stringify({
        type: 'progress',
        data: { prompt_id: jobId, node: nodeId, value, max }
      })
    )
  }

  /**
   * Mock the history API to return the given job as completed,
   * then send execution_success and a status event to trigger
   * the history refresh that resolves in-progress items.
   */
  async completeWithHistory(
    jobId: string,
    nodeId: string,
    filename: string
  ): Promise<void> {
    const now = Date.now() / 1000
    const job: RawJobListItem = {
      id: jobId,
      status: 'completed',
      create_time: now,
      execution_start_time: now,
      execution_end_time: now + 1,
      priority: 0,
      outputs_count: 1,
      preview_output: {
        filename,
        subfolder: '',
        type: 'output',
        nodeId,
        mediaType: 'images'
      }
    }

    await this.page.route(/\/api\/jobs\?status=completed/, async (route) => {
      await route.fulfill({
        status: 200,
        contentType: 'application/json',
        body: JSON.stringify({
          jobs: [job],
          pagination: { offset: 0, limit: 200, total: 1, has_more: false }
        })
      })
    })

    this.executionSuccess(jobId)
    // Trigger queue/history refresh
    this.status(0)
  }

  /** Send `status` WS event to update queue count. */
  status(queueRemaining: number): void {
    this.ws.send(
      JSON.stringify({
        type: 'status',
        data: { status: { exec_info: { queue_remaining: queueRemaining } } }
      })
    )
  }
}