ComfyUI_frontend/tests-ui/platform/assets/utils/createModelNodeFromAsset.test.ts

import { beforeEach, describe, expect, it, vi } from 'vitest'
import { type Raw, markRaw } from 'vue'

import {
  type LGraphNode,
  LiteGraph,
  type Subgraph
} from '@/lib/litegraph/src/litegraph'
import type { AssetItem } from '@/platform/assets/schemas/assetSchema'
import { createModelNodeFromAsset } from '@/platform/assets/utils/createModelNodeFromAsset'
import { useWorkflowStore } from '@/platform/workflow/management/stores/workflowStore'
import { app } from '@/scripts/app'
import { useLitegraphService } from '@/services/litegraphService'
import { useModelToNodeStore } from '@/stores/modelToNodeStore'

// Mock dependencies
vi.mock('@/scripts/api', () => ({
  api: {
    apiURL: vi.fn((path: string) => `http://localhost:8188${path}`)
  }
}))
vi.mock('@/stores/modelToNodeStore', async (importOriginal) => {
  const actual =
    await importOriginal<typeof import('@/stores/modelToNodeStore')>()
  return {
    ...actual,
    useModelToNodeStore: vi.fn()
  }
})
vi.mock(
  '@/platform/workflow/management/stores/workflowStore',
  async (importOriginal) => {
    const actual =
      await importOriginal<
        typeof import('@/platform/workflow/management/stores/workflowStore')
      >()
    return {
      ...actual,
      useWorkflowStore: vi.fn()
    }
  }
)
vi.mock('@/services/litegraphService', async (importOriginal) => {
  const actual =
    await importOriginal<typeof import('@/services/litegraphService')>()
  return {
    ...actual,
    useLitegraphService: vi.fn()
  }
})
vi.mock('@/lib/litegraph/src/litegraph', async (importOriginal) => {
  const actual =
    await importOriginal<typeof import('@/lib/litegraph/src/litegraph')>()
  return {
    ...actual,
    LiteGraph: {
      ...actual.LiteGraph,
      createNode: vi.fn()
    }
  }
})
vi.mock('@/scripts/app', () => ({
  app: {
    canvas: {
      graph: {
        add: vi.fn()
      }
    }
  }
}))
function createMockAsset(overrides: Partial<AssetItem> = {}): AssetItem {
  return {
    id: 'asset-123',
    name: 'test-model.safetensors',
    size: 1024,
    created_at: '2025-10-01T00:00:00Z',
    tags: ['models', 'checkpoints'],
    user_metadata: {
      filename: 'models/checkpoints/test-model.safetensors'
    },
    ...overrides
  }
}
async function createMockNode(overrides?: {
  widgetName?: string
  widgetValue?: string
  hasWidgets?: boolean
}): Promise<LGraphNode> {
  const {
    widgetName = 'ckpt_name',
    widgetValue = '',
    hasWidgets = true
  } = overrides || {}

  const { LGraphNode: ActualLGraphNode } = await vi.importActual<
    typeof import('@/lib/litegraph/src/litegraph')
  >('@/lib/litegraph/src/litegraph')

  if (!hasWidgets) {
    return Object.create(ActualLGraphNode.prototype)
  }

  type Widget = NonNullable<LGraphNode['widgets']>[number]
  const widget: Pick<Widget, 'name' | 'value' | 'type' | 'options' | 'y'> = {
    name: widgetName,
    value: widgetValue,
    type: 'string',
    options: {},
    y: 0
  }

  return Object.create(ActualLGraphNode.prototype, {
    widgets: { value: [widget], writable: true }
  })
}
function createMockNodeProvider() {
  return {
    nodeDef: {
      name: 'CheckpointLoaderSimple',
      display_name: 'Load Checkpoint'
    },
    key: 'ckpt_name'
  }
}
/**
 * Configures all mocked dependencies with sensible defaults.
 * Uses semantic parameters for clearer test intent.
 * For error paths or edge cases, pass null values or specific overrides.
 */
async function setupMocks(
  overrides: {
    nodeProvider?: ReturnType<typeof createMockNodeProvider> | null
    canvasCenter?: [number, number]
    activeSubgraph?: Raw<Subgraph>
    createdNode?: Awaited<ReturnType<typeof createMockNode>> | null
  } = {}
) {
  const {
    nodeProvider = createMockNodeProvider(),
    canvasCenter = [100, 200],
    activeSubgraph = undefined,
    createdNode = await createMockNode()
  } = overrides

  vi.mocked(useModelToNodeStore).mockReturnValue({
    ...useModelToNodeStore(),
    getNodeProvider: vi.fn().mockReturnValue(nodeProvider)
  })

  vi.mocked(useLitegraphService).mockReturnValue({
    ...useLitegraphService(),
    getCanvasCenter: vi.fn().mockReturnValue(canvasCenter)
  })

  vi.mocked(useWorkflowStore).mockReturnValue({
    ...useWorkflowStore(),
    activeSubgraph,
    isSubgraphActive: !!activeSubgraph
  })
  vi.mocked(LiteGraph.createNode).mockReturnValue(createdNode)
}
describe('createModelNodeFromAsset', () => {
  beforeEach(() => {
    vi.clearAllMocks()
    vi.spyOn(console, 'warn').mockImplementation(() => {})
  })
  describe('when creating nodes from valid assets', () => {
    it('should create the appropriate loader node for the asset category', async () => {
      const asset = createMockAsset()
      await setupMocks()
      const result = createModelNodeFromAsset(asset)
      expect(result.success).toBe(true)
      if (result.success) {
        expect(
          vi.mocked(useModelToNodeStore)().getNodeProvider
        ).toHaveBeenCalledWith('checkpoints')
        expect(LiteGraph.createNode).toHaveBeenCalledWith(
          'CheckpointLoaderSimple',
          'Load Checkpoint',
          { pos: [100, 200] }
        )
      }
    })
    it('should place node at canvas center by default', async () => {
      const asset = createMockAsset()
      await setupMocks({
        canvasCenter: [150, 250]
      })
      const result = createModelNodeFromAsset(asset)
      expect(result.success).toBe(true)
      expect(
        vi.mocked(useLitegraphService)().getCanvasCenter
      ).toHaveBeenCalled()
      expect(LiteGraph.createNode).toHaveBeenCalledWith(
        'CheckpointLoaderSimple',
        'Load Checkpoint',
        { pos: [150, 250] }
      )
    })
    it('should place node at specified position when position is provided', async () => {
      const asset = createMockAsset()
      await setupMocks()
      const result = createModelNodeFromAsset(asset, { position: [300, 400] })
      expect(result.success).toBe(true)
      expect(
        vi.mocked(useLitegraphService)().getCanvasCenter
      ).not.toHaveBeenCalled()
      expect(LiteGraph.createNode).toHaveBeenCalledWith(
        'CheckpointLoaderSimple',
        'Load Checkpoint',
        { pos: [300, 400] }
      )
    })
    it('should populate the loader widget with the asset file path', async () => {
      const asset = createMockAsset()
      const mockNode = await createMockNode()
      await setupMocks({ createdNode: mockNode })
      const result = createModelNodeFromAsset(asset)
      expect(result.success).toBe(true)
      expect(mockNode.widgets?.[0].value).toBe(
        'models/checkpoints/test-model.safetensors'
      )
    })
    it('should add node to root graph when no subgraph is active', async () => {
      const asset = createMockAsset()
      const mockNode = await createMockNode()
      await setupMocks({ createdNode: mockNode })
      const result = createModelNodeFromAsset(asset)
      expect(result.success).toBe(true)
      expect(vi.mocked(app).canvas.graph!.add).toHaveBeenCalledWith(mockNode)
    })
    it('should add node to active subgraph when present', async () => {
      const asset = createMockAsset()
      const mockNode = await createMockNode()
      const { Subgraph } = await vi.importActual<
        typeof import('@/lib/litegraph/src/litegraph')
      >('@/lib/litegraph/src/litegraph')
      const mockSubgraph = markRaw(
        Object.create(Subgraph.prototype, {
          add: { value: vi.fn() }
        })
      )
      await setupMocks({
        createdNode: mockNode,
        activeSubgraph: mockSubgraph
      })
      const result = createModelNodeFromAsset(asset)
      expect(result.success).toBe(true)
      expect(mockSubgraph.add).toHaveBeenCalledWith(mockNode)
      expect(vi.mocked(app).canvas.graph!.add).not.toHaveBeenCalled()
    })
  })
  describe('when asset data is incomplete or invalid', () => {
    beforeEach(() => {
      vi.spyOn(console, 'error').mockImplementation(() => {})
    })
    it.each([
      {
        case: 'missing user_metadata',
        overrides: { user_metadata: undefined },
        expectedCode: 'INVALID_ASSET' as const,
        errorPattern: /missing required user_metadata/
      },
      {
        case: 'missing filename property',
        overrides: { user_metadata: {} },
        expectedCode: 'INVALID_ASSET' as const,
        errorPattern:
          /Invalid filename.*expected non-empty string, got undefined/
      },
      {
        case: 'non-string filename',
        overrides: { user_metadata: { filename: 123 } },
        expectedCode: 'INVALID_ASSET' as const,
        errorPattern: /Invalid filename.*expected non-empty string, got number/
      },
      {
        case: 'empty filename',
        overrides: { user_metadata: { filename: '' } },
        expectedCode: 'INVALID_ASSET' as const,
        errorPattern: /Invalid filename.*expected non-empty string/
      }
    ])(
      'should fail when asset has $case',
      ({ overrides, expectedCode, errorPattern }) => {
        const asset = createMockAsset(overrides)
        const result = createModelNodeFromAsset(asset)
        expect(result.success).toBe(false)
        if (!result.success) {
          expect(result.error.code).toBe(expectedCode)
          expect(result.error.message).toMatch(errorPattern)
          expect(result.error.assetId).toBe('asset-123')
        }
      }
    )
    it.each([
      {
        case: 'no tags',
        overrides: { tags: undefined },
        expectedCode: 'INVALID_ASSET' as const,
        errorMessage: 'Asset has no tags defined'
      },
      {
        case: 'only excluded tags',
        overrides: { tags: ['models', 'missing'] },
        expectedCode: 'INVALID_ASSET' as const,
        errorMessage: 'Asset has no valid category tag'
      },
      {
        case: 'only the models tag',
        overrides: { tags: ['models'] },
        expectedCode: 'INVALID_ASSET' as const,
        errorMessage: 'Asset has no valid category tag'
      }
    ])(
      'should fail when asset has $case',
      ({ overrides, expectedCode, errorMessage }) => {
        const asset = createMockAsset(overrides)
        const result = createModelNodeFromAsset(asset)
        expect(result.success).toBe(false)
        if (!result.success) {
          expect(result.error.code).toBe(expectedCode)
          expect(result.error.message).toBe(errorMessage)
        }
      }
    )
  })
  describe('when system resources are unavailable', () => {
    beforeEach(() => {
      vi.spyOn(console, 'error').mockImplementation(() => {})
    })
    it('should fail when no provider registered for category', async () => {
      const asset = createMockAsset()
      await setupMocks({ nodeProvider: null })
      const result = createModelNodeFromAsset(asset)
      expect(result.success).toBe(false)
      if (!result.success) {
        expect(result.error.code).toBe('NO_PROVIDER')
        expect(result.error.message).toContain('checkpoints')
        expect(result.error.details?.category).toBe('checkpoints')
      }
    })
    it('should fail when node creation fails', async () => {
      const asset = createMockAsset()
      await setupMocks()
      vi.mocked(LiteGraph.createNode).mockReturnValue(null)
      const result = createModelNodeFromAsset(asset)
      expect(result.success).toBe(false)
      if (!result.success) {
        expect(result.error.code).toBe('NODE_CREATION_FAILED')
        expect(result.error.message).toContain('CheckpointLoaderSimple')
      }
    })
    it('should fail when widget is missing from node', async () => {
      const asset = createMockAsset()
      const mockNode = await createMockNode({ widgetName: 'wrong_widget' })
      await setupMocks({ createdNode: mockNode })
      const result = createModelNodeFromAsset(asset)
      expect(result.success).toBe(false)
      if (!result.success) {
        expect(result.error.code).toBe('MISSING_WIDGET')
        expect(result.error.message).toContain('ckpt_name')
        expect(result.error.message).toContain('CheckpointLoaderSimple')
        expect(result.error.details?.widgetName).toBe('ckpt_name')
      }
    })
    it('should fail when node has no widgets array', async () => {
      const asset = createMockAsset()
      const mockNode = await createMockNode({ hasWidgets: false })
      await setupMocks({ createdNode: mockNode })
      const result = createModelNodeFromAsset(asset)
      expect(result.success).toBe(false)
      if (!result.success) {
        expect(result.error.code).toBe('MISSING_WIDGET')
        expect(result.error.message).toContain('ckpt_name not found')
      }
    })
    it('should not add node to graph when widget validation fails', async () => {
      const asset = createMockAsset()
      const mockNode = await createMockNode({ hasWidgets: false })
      await setupMocks({ createdNode: mockNode })
      createModelNodeFromAsset(asset)
      expect(vi.mocked(app).canvas.graph!.add).not.toHaveBeenCalled()
    })
  })
  describe('when graph is null', () => {
    beforeEach(() => {
      vi.spyOn(console, 'error').mockImplementation(() => {})
      vi.mocked(app).canvas.graph = null
    })
    it('should fail when no graph is available', async () => {
      const asset = createMockAsset()
      const mockNode = await createMockNode()
      await setupMocks({ createdNode: mockNode })
      const result = createModelNodeFromAsset(asset)
      expect(result.success).toBe(false)
      if (!result.success) {
        expect(result.error.code).toBe('NO_GRAPH')
        expect(result.error.message).toBe('No active graph available')
      }
    })
  })
})