deepseek-harness/packages/extensions/tool-lab/tests/tool-lab.spec.ts

/**
 * Unit tests for the pure helpers of `@deepseek-ai/dsh-tool-lab`. These
 * assertions do not touch the network or the runtime; they check the workflow
 * and upload builders only.
 * @module @deepseek-ai/dsh-tool-lab/tests
 */

import { describe, expect, it } from 'vitest'
import { buildWorkflow } from '../src/comfy.ts'
import { buildPdfUpload } from '../src/docling.ts'
import { buildAudioUpload } from '../src/whish.ts'

/** One ComfyUI workflow node as seen by the assertions. */
interface WorkflowNode {
  class_type: string
  inputs: Record<string, unknown>
}

/** Reinterpret the workflow builder's JSON-serializable output for assertions. */
function wf(prompt: string, overrides: Record<string, unknown> = {}): Record<string, WorkflowNode> {
  return buildWorkflow({ prompt, ...(overrides as object) }) as Record<string, WorkflowNode>
}

describe('buildWorkflow', () => {
  it('renders a 28-step Juggernaut workflow with defaults', () => {
    const flow = wf('a red fox')
    expect(flow['4'].inputs.ckpt_name).toBe('Juggernaut-XL_v9_RunDiffusionPhoto_v2.safetensors')
    expect(flow['9'].class_type).toBe('SaveImage')
    expect(flow['9'].inputs.images).toEqual(['10', 0])
    expect(flow['10'].class_type).toBe('ImageUpscaleWithModel')
    expect(flow['11'].inputs.model_name).toBe('RealESRGAN_x4plus.safetensors')
    expect(flow['3'].inputs).toMatchObject({ steps: 28, cfg: 4.0, sampler_name: 'dpmpp_2m_sde', scheduler: 'karras' })
  })

  it('skips upscaling when upscale is false', () => {
    const flow = wf('p', { upscale: false })
    expect(flow['9'].inputs.images).toEqual(['8', 0])
    expect(flow['10']).toBeUndefined()
    expect(flow['11']).toBeUndefined()
  })

  it('honors overrides and seed', () => {
    const flow = wf('p', { width: 768, height: 640, steps: 8, seed: 99 })
    expect(flow['5'].inputs).toMatchObject({ width: 768, height: 640 })
    expect(flow['3'].inputs).toMatchObject({ seed: 99, steps: 8 })
    expect(flow['6'].inputs.text).toBe('p')
  })
})

describe('upload builders', () => {
  it('builds a PDF upload carrying files and options parts', () => {
    const fd = buildPdfUpload(new Uint8Array([1, 2, 3]), 'scan.pdf')
    expect(fd).toBeDefined()
  })

  it('builds an audio upload', () => {
    const fd = buildAudioUpload(new Uint8Array([1, 2]), 'clip.mp3')
    expect(fd).toBeDefined()
  })
})