ComfyUI_frontend/browser_tests/fixtures/helpers/Preview3DPipelineFixture.ts

import { comfyPageFixture } from '@e2e/fixtures/ComfyPage'
import type { ComfyPage } from '@e2e/fixtures/ComfyPage'
import { Load3DHelper } from '@e2e/tests/load3d/Load3DHelper'

export async function waitForAppBootstrapped(
  comfyPage: ComfyPage
): Promise<void> {
  await comfyPage.page.waitForFunction(() =>
    Boolean(window.app?.extensionManager)
  )
  await comfyPage.page.waitForSelector('.p-blockui-mask', { state: 'hidden' })
  await comfyPage.nextFrame()
}

export function vecClose(
  a: { x: number; y: number; z: number },
  b: { x: number; y: number; z: number },
  eps: number
): boolean {
  return (
    Math.abs(a.x - b.x) < eps &&
    Math.abs(a.y - b.y) < eps &&
    Math.abs(a.z - b.z) < eps
  )
}

export function cameraStatesClose(
  a: unknown,
  b: unknown,
  eps: number
): boolean {
  if (
    a === null ||
    b === null ||
    typeof a !== 'object' ||
    typeof b !== 'object'
  )
    return false
  const ca = a as {
    position?: { x: number; y: number; z: number }
    target?: { x: number; y: number; z: number }
    zoom?: number
    cameraType?: string
  }
  const cb = b as typeof ca
  if (
    !ca.position ||
    !ca.target ||
    !cb.position ||
    !cb.target ||
    ca.cameraType !== cb.cameraType
  )
    return false
  if (Math.abs((ca.zoom ?? 0) - (cb.zoom ?? 0)) > eps) return false
  return (
    vecClose(ca.position, cb.position, eps) &&
    vecClose(ca.target, cb.target, eps)
  )
}

export class Preview3DPipelineContext {
  static readonly loadNodeId = '1'
  static readonly previewNodeId = '2'

  readonly load3d: Load3DHelper
  readonly preview3d: Load3DHelper

  constructor(readonly comfyPage: ComfyPage) {
    this.load3d = new Load3DHelper(
      comfyPage.vueNodes.getNodeLocator(Preview3DPipelineContext.loadNodeId)
    )
    this.preview3d = new Load3DHelper(
      comfyPage.vueNodes.getNodeLocator(Preview3DPipelineContext.previewNodeId)
    )
  }

  async getModelFileWidgetValue(nodeId: string): Promise<string> {
    return this.comfyPage.page.evaluate((id) => {
      const n = window.app!.graph.getNodeById(Number(id))
      const w = n?.widgets?.find((x) => x.name === 'model_file')
      return typeof w?.value === 'string' ? w.value : ''
    }, nodeId)
  }

  async getLastTimeModelFile(nodeId: string): Promise<string> {
    return this.comfyPage.page.evaluate((id) => {
      const n = window.app!.graph.getNodeById(Number(id))
      const v = n?.properties?.['Last Time Model File']
      return typeof v === 'string' ? v : ''
    }, nodeId)
  }

  async getCameraStateFromProperties(nodeId: string): Promise<unknown> {
    return this.comfyPage.page.evaluate((id) => {
      const n = window.app!.graph.getNodeById(Number(id))
      const cfg = n?.properties?.['Camera Config'] as
        | { state?: unknown }
        | undefined
      return cfg?.state ?? null
    }, nodeId)
  }
}

export const preview3dPipelineTest = comfyPageFixture.extend<{
  preview3dPipeline: Preview3DPipelineContext
}>({
  preview3dPipeline: async ({ comfyPage }, use) => {
    await comfyPage.settings.setSetting('Comfy.VueNodes.Enabled', true)
    await comfyPage.settings.setSetting(
      'Comfy.Workflow.WorkflowTabsPosition',
      'Sidebar'
    )
    await comfyPage.workflow.loadWorkflow('3d/preview3d_pipeline')
    await comfyPage.vueNodes.waitForNodes()

    const pipeline = new Preview3DPipelineContext(comfyPage)
    await use(pipeline)

    await comfyPage.workflow.setupWorkflowsDirectory({})
  }
})