53 lines
2.2 KiB
JavaScript
53 lines
2.2 KiB
JavaScript
import assert from 'node:assert/strict'
|
|
import { createCaptionHost, validateCaptionHostRequest } from '../scripts/caption-host.mjs'
|
|
import { CAPTION_STYLES, captionStylePrompt, DELTA_LOCKED_LOOK, validateCaptionBody } from '../shared/caption.mjs'
|
|
import { mkdtempSync, writeFileSync, mkdirSync, existsSync } from 'node:fs'
|
|
import { join } from 'node:path'
|
|
import { tmpdir } from 'node:os'
|
|
import { EventEmitter } from 'node:events'
|
|
|
|
assert.deepEqual(CAPTION_STYLES, ['descriptive', 'klein_prompt', 'delta', 'tags'])
|
|
assert.match(captionStylePrompt('delta'), new RegExp(DELTA_LOCKED_LOOK.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')))
|
|
assert.equal(validateCaptionBody({ folderId: 'f1', stillId: 's1', captionStyle: 'tags' }).captionStyle, 'tags')
|
|
assert.throws(() => validateCaptionBody({ folderId: 'f1' }), /still/)
|
|
assert.throws(() => validateCaptionHostRequest({ id: 'short', style: 'descriptive' }), /Invalid/)
|
|
|
|
const root = mkdtempSync(join(tmpdir(), 'caption-test-'))
|
|
const models = join(root, 'models')
|
|
mkdirSync(models)
|
|
writeFileSync(join(models, 'Qwen2.5-VL-7B-NSFW-Caption-V4.Q5_K_M.gguf'), 'x')
|
|
writeFileSync(join(models, 'Qwen2.5-VL-7B-NSFW-Caption-V4.mmproj-f16.gguf'), 'x')
|
|
writeFileSync(join(root, 'llama-server.exe'), 'x')
|
|
|
|
let prepared = 0
|
|
const host = createCaptionHost({
|
|
modelsDir: models,
|
|
llamaServer: join(root, 'llama-server.exe'),
|
|
dataDir: join(root, 'jobs'),
|
|
leaseValid: () => true,
|
|
prepare: async () => { prepared += 1 },
|
|
serverWaitMs: 800,
|
|
spawnProcess: () => {
|
|
const child = new EventEmitter()
|
|
child.pid = 4242
|
|
child.stdout = new EventEmitter()
|
|
child.stderr = new EventEmitter()
|
|
child.kill = () => child.emit('close', 0)
|
|
setTimeout(() => child.emit('close', 0), 50)
|
|
return child
|
|
},
|
|
fetchImpl: async () => ({ ok: false, status: 503, text: async () => 'down' })
|
|
})
|
|
|
|
assert.equal(host.configured(), true)
|
|
assert.equal(host.busy(), false)
|
|
|
|
const id = 'caption-test-abcdefghijkl'
|
|
mkdirSync(join(root, 'jobs', id), { recursive: true })
|
|
writeFileSync(join(root, 'jobs', id, 'input.png'), Buffer.alloc(64, 1))
|
|
|
|
await assert.rejects(() => host.start({ id, style: 'descriptive' }, 'lease'), /llama-server|Caption/)
|
|
assert.equal(prepared, 1)
|
|
assert.ok(existsSync(join(root, 'jobs', id, 'status.json')))
|
|
console.log('caption tests ok')
|