ComfyUI_frontend/browser_tests/fixtures/helpers/Preview3DPipelineFixture.ts

import { expect } from '@playwright/test'

import { comfyPageFixture } from '@e2e/fixtures/ComfyPage'
import type { ComfyPage } from '@e2e/fixtures/ComfyPage'
import { assetPath } from '@e2e/fixtures/utils/paths'
import type { Preview3dE2eCameraState as CameraState } from '@e2e/fixtures/utils/preview3dCameraState'
import {
  isPreview3dE2eCameraState as isCameraState,
  preview3dCameraStatesDiffer as cameraStatesDiffer
} from '@e2e/fixtures/utils/preview3dCameraState'
import { Load3DHelper } from '@e2e/tests/load3d/Load3DHelper'

export type { CameraState }
export { isCameraState }

export async function backendHasPreview3DNodes(
  comfyPage: ComfyPage
): Promise<boolean> {
  const resp = await comfyPage.request.get(`${comfyPage.url}/object_info`, {
    failOnStatusCode: false
  })
  if (!resp.ok()) return false
  const data: unknown = await resp.json()
  return (
    typeof data === 'object' &&
    data !== null &&
    'Load3D' in data &&
    'Preview3D' in data
  )
}

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> {
    const ref = await this.comfyPage.nodeOps.getNodeRefById(nodeId)
    const v = await ref.getWidgetValueByName('model_file')
    return typeof v === 'string' ? v : ''
  }

  async getLastTimeModelFile(nodeId: string): Promise<string> {
    const ref = await this.comfyPage.nodeOps.getNodeRefById(nodeId)
    const v = await ref.getStoredPropertyValue('Last Time Model File')
    return typeof v === 'string' ? v : ''
  }

  async getCameraStateFromProperties(nodeId: string): Promise<unknown> {
    const ref = await this.comfyPage.nodeOps.getNodeRefById(nodeId)
    const cfg = await ref.getStoredPropertyValue('Camera Config')
    if (cfg === null || typeof cfg !== 'object') return null
    const state = Reflect.get(cfg, 'state')
    return state ?? null
  }
}

export async function seedLoad3dWithCubeObj(
  pipeline: Preview3DPipelineContext
): Promise<void> {
  const { comfyPage, load3d } = pipeline
  const fileChooserPromise = comfyPage.page.waitForEvent('filechooser')
  await load3d.getUploadButton('upload 3d model').click()
  const fileChooser = await fileChooserPromise
  await fileChooser.setFiles(assetPath('cube.obj'))

  await expect
    .poll(() =>
      pipeline.getModelFileWidgetValue(Preview3DPipelineContext.loadNodeId)
    )
    .toContain('cube.obj')

  await load3d.waitForModelLoaded()
  await comfyPage.nextFrame()
}

export async function setNonDefaultLoad3dCameraState(
  pipeline: Preview3DPipelineContext
): Promise<void> {
  const { comfyPage, load3d } = pipeline
  const initialCamera = await pipeline.getCameraStateFromProperties(
    Preview3DPipelineContext.loadNodeId
  )
  const box = await load3d.canvas.boundingBox()
  expect(box, 'Load3D canvas bounding box should exist').not.toBeNull()

  const w = box!.width
  const h = box!.height
  await load3d.canvas.dragTo(load3d.canvas, {
    sourcePosition: { x: w / 2, y: h / 2 },
    targetPosition: { x: w / 2 + 80, y: h / 2 + 20 },
    steps: 10,
    force: true
  })
  await comfyPage.nextFrame()

  await expect
    .poll(
      async () => {
        const current = await pipeline.getCameraStateFromProperties(
          Preview3DPipelineContext.loadNodeId
        )
        if (current === null) return false
        if (initialCamera === null) return true
        return cameraStatesDiffer(current, initialCamera, 1e-4)
      },
      {
        timeout: 10_000,
        message:
          'Load3D camera state should change after orbit drag (see cameraStatesDiffer)'
      }
    )
    .toBe(true)
}

export async function nudgePreview3dCameraIntoProperties(
  pipeline: Preview3DPipelineContext
): Promise<void> {
  const { comfyPage, preview3d } = pipeline
  const box = await preview3d.canvas.boundingBox()
  expect(box, 'Preview3D canvas bounding box should exist').not.toBeNull()

  const w = box!.width
  const h = box!.height
  await preview3d.canvas.dragTo(preview3d.canvas, {
    sourcePosition: { x: w / 2, y: h / 2 },
    targetPosition: { x: w / 2 - 60, y: h / 2 + 20 },
    steps: 10,
    force: true
  })
  await comfyPage.nextFrame()
}

export async function alignPreview3dWorkflowUiSettings(
  pipeline: Preview3DPipelineContext
): Promise<void> {
  await pipeline.comfyPage.settings.setSetting('Comfy.VueNodes.Enabled', true)
  await pipeline.comfyPage.settings.setSetting(
    'Comfy.Workflow.WorkflowTabsPosition',
    'Sidebar'
  )
}

export const preview3dPipelineTest = comfyPageFixture.extend<{
  preview3dPipeline: Preview3DPipelineContext
}>({
  preview3dPipeline: async ({ comfyPage }, use, testInfo) => {
    await comfyPage.settings.setSetting('Comfy.VueNodes.Enabled', true)
    await comfyPage.settings.setSetting(
      'Comfy.Workflow.WorkflowTabsPosition',
      'Sidebar'
    )

    const hasPreview3dNodes = await backendHasPreview3DNodes(comfyPage)
    if (!hasPreview3dNodes) {
      testInfo.skip(
        true,
        'Requires ComfyUI backend with Load3D and Preview3D nodes (object_info)'
      )
      await use(new Preview3DPipelineContext(comfyPage))
      await comfyPage.workflow.setupWorkflowsDirectory({})
      return
    }

    await comfyPage.workflow.loadWorkflow('3d/preview3d_pipeline')
    await comfyPage.vueNodes.waitForNodes()

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

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