90 lines
4.8 KiB
TypeScript
90 lines
4.8 KiB
TypeScript
import { afterEach, describe, expect, it, vi } from 'vitest'
|
|
import { runComfyImage, resolveImageArgs, buildWorkflow } from '@deepseek-ai/dsh-tool-lab/src/comfy.ts'
|
|
import { runDoclingOcr } from '@deepseek-ai/dsh-tool-lab/src/docling.ts'
|
|
import { runWhishTranscribe } from '@deepseek-ai/dsh-tool-lab/src/whish.ts'
|
|
|
|
const pngBytes = new Uint8Array([137, 80, 78, 71, 13, 10, 26, 10, 0, 0, 0, 13])
|
|
|
|
function jsonResponse(body: unknown, { status = 200, headers }: { status?: number; headers?: Record<string, string> } = {}): Response {
|
|
return new Response(JSON.stringify(body), { status, headers: { 'content-type': 'application/json', ...headers } })
|
|
}
|
|
|
|
afterEach(() => { vi.unstubAllGlobals() })
|
|
|
|
describe('lab_generate_image (ComfyUI)', () => {
|
|
it('queues the workflow and returns the served PNG URL after polling history', async () => {
|
|
const calls: string[] = []
|
|
vi.stubGlobal('fetch', vi.fn(async (input: RequestInfo | URL) => {
|
|
const url = String(input)
|
|
calls.push(url)
|
|
if (url.endsWith('/prompt')) return jsonResponse({ prompt_id: 'p-1', node_errors: {} })
|
|
if (url.includes('/history/p-1')) return jsonResponse({ 'p-1': { outputs: { '9': { images: [{ filename: 'dsh_00001_.png' }] } } } })
|
|
throw new Error(`unexpected fetch ${url}`)
|
|
}))
|
|
|
|
const out = await runComfyImage('http://192.168.31.240:8188', resolveImageArgs({ prompt: 'a cat' }), undefined, 30_000)
|
|
expect(out).toBe('http://192.168.31.240:8188/view?filename=dsh_00001_.png&subfolder=&type=output')
|
|
expect(calls.some(u => u.endsWith('/prompt'))).toBe(true)
|
|
})
|
|
|
|
it('rejects a workflow with non-empty node_errors', async () => {
|
|
vi.stubGlobal('fetch', vi.fn(async () => jsonResponse({ prompt_id: 'p-2', node_errors: { '3': ['bad'] } })))
|
|
await expect(runComfyImage('http://x', resolveImageArgs({ prompt: 'x' }), undefined, 30_000))
|
|
.rejects.toThrow(/workflow rejected/)
|
|
})
|
|
|
|
it('rejects non-positive integer dimensions and steps', () => {
|
|
expect(() => resolveImageArgs({ prompt: 'x', width: 0 })).toThrow(/width must be a positive integer/)
|
|
expect(() => resolveImageArgs({ prompt: 'x', steps: -1 })).toThrow(/steps must be a positive integer/)
|
|
})
|
|
|
|
it('uses the documented SDXL-turbo defaults in the workflow', () => {
|
|
const wf = buildWorkflow({ prompt: 'p', negative: 'negative prompt', width: 512, height: 512, steps: 4, seed: 123, model: 'sdxl_turbo.safetensors', filenamePrefix: 'dsh' })
|
|
const sampler = wf['3'] as { inputs: Record<string, unknown> }
|
|
expect(sampler.inputs.cfg).toBe(1.0)
|
|
expect(sampler.inputs.sampler_name).toBe('euler')
|
|
expect(sampler.inputs.scheduler).toBe('normal')
|
|
expect(sampler.inputs.steps).toBe(4)
|
|
})
|
|
})
|
|
|
|
describe('lab_ocr_pdf (Docling)', () => {
|
|
it('uploads a PDF and returns document.md_content after polling', async () => {
|
|
vi.stubGlobal('fetch', vi.fn(async (input: RequestInfo | URL) => {
|
|
const url = String(input)
|
|
if (url.includes('/v1/convert/file/async')) return jsonResponse({ task_id: 't-1', task_status: 'pending' })
|
|
if (url.includes('/v1/status/poll/t-1')) return jsonResponse({ task_status: 'success' })
|
|
if (url.includes('/v1/result/t-1')) return jsonResponse({ status: 'success', errors: [], document: { filename: 'a.pdf', md_content: '# OCRed' } })
|
|
throw new Error(`unexpected ${url}`)
|
|
}))
|
|
const out = await runDoclingOcr('http://192.168.31.159:5001', new TextEncoder().encode('%PDF-1.4'), 'a.pdf', undefined, 30_000)
|
|
expect(out).toBe('# OCRed')
|
|
expect(globalThis.fetch).toHaveBeenCalled()
|
|
})
|
|
|
|
it('rejects an empty task id', async () => {
|
|
vi.stubGlobal('fetch', vi.fn(async () => jsonResponse({ task_id: undefined })))
|
|
await expect(runDoclingOcr('http://x', new Uint8Array(), 'a.pdf', undefined, 30_000)).rejects.toThrow(/no task_id/)
|
|
})
|
|
})
|
|
|
|
describe('lab_transcribe_audio (Whishper)', () => {
|
|
it('polls until result.text is non-empty', async () => {
|
|
let polls = 0
|
|
vi.stubGlobal('fetch', vi.fn(async (input: RequestInfo | URL) => {
|
|
const url = String(input)
|
|
if (url.endsWith('/api/transcriptions')) return jsonResponse({ id: 'w-1', status: -1 })
|
|
polls += 1
|
|
if (polls === 1) return jsonResponse({ id: 'w-1', status: -1, result: { text: '' } })
|
|
return jsonResponse({ id: 'w-1', status: 0, result: { text: 'hello world', language: 'en', duration: 1.2 } })
|
|
}))
|
|
const out = await runWhishTranscribe('http://192.168.31.159:8082', new TextEncoder().encode('wav'), 'a.wav', undefined, undefined, undefined, 30_000)
|
|
expect(out).toBe('hello world')
|
|
})
|
|
|
|
it('rejects an upload with no id', async () => {
|
|
vi.stubGlobal('fetch', vi.fn(async () => jsonResponse({ id: undefined })))
|
|
await expect(runWhishTranscribe('http://x', new Uint8Array(), 'a.wav', undefined, undefined, undefined, 30_000))
|
|
.rejects.toThrow(/no id/)
|
|
})
|
|
}) |