feat: initial release v0.3.0
This commit is contained in:
saturn
240
tests/unit/worker/reference-to-character.test.ts
Normal file
240
tests/unit/worker/reference-to-character.test.ts
Normal file
@@ -0,0 +1,240 @@
|
||||
import type { Job } from 'bullmq'
|
||||
import { beforeEach, describe, expect, it, vi } from 'vitest'
|
||||
import { CHARACTER_PROMPT_SUFFIX, CHARACTER_IMAGE_BANANA_RATIO } from '@/lib/constants'
|
||||
import { TASK_TYPE, type TaskJobData, type TaskType } from '@/lib/task/types'
|
||||
|
||||
const sharpMock = vi.hoisted(() =>
|
||||
vi.fn(() => {
|
||||
const chain = {
|
||||
metadata: vi.fn(async () => ({ width: 2160, height: 2160 })),
|
||||
extend: vi.fn(() => chain),
|
||||
composite: vi.fn(() => chain),
|
||||
jpeg: vi.fn(() => chain),
|
||||
toBuffer: vi.fn(async () => Buffer.from('processed-image')),
|
||||
}
|
||||
return chain
|
||||
}),
|
||||
)
|
||||
|
||||
const generatorApiMock = vi.hoisted(() => ({
|
||||
generateImage: vi.fn<(userId: string, modelId: string, prompt: string, options?: Record<string, unknown>) => Promise<{
|
||||
success: boolean
|
||||
imageUrl: string
|
||||
async: boolean
|
||||
}>>(async () => ({
|
||||
success: true,
|
||||
imageUrl: 'https://example.com/generated.jpg',
|
||||
async: false,
|
||||
})),
|
||||
}))
|
||||
|
||||
const asyncSubmitMock = vi.hoisted(() => ({
|
||||
queryFalStatus: vi.fn(async () => ({ completed: false, failed: false, resultUrl: null })),
|
||||
}))
|
||||
|
||||
const arkApiMock = vi.hoisted(() => ({
|
||||
fetchWithTimeoutAndRetry: vi.fn(async () => ({
|
||||
arrayBuffer: async () => new Uint8Array([1, 2, 3]).buffer,
|
||||
})),
|
||||
}))
|
||||
|
||||
const apiConfigMock = vi.hoisted(() => ({
|
||||
getProviderConfig: vi.fn(async () => ({ apiKey: 'fal-key' })),
|
||||
}))
|
||||
|
||||
const configServiceMock = vi.hoisted(() => ({
|
||||
getUserModelConfig: vi.fn(async () => ({
|
||||
characterModel: 'character-model-1',
|
||||
analysisModel: 'analysis-model-1',
|
||||
})),
|
||||
}))
|
||||
|
||||
const llmClientMock = vi.hoisted(() => ({
|
||||
chatCompletionWithVision: vi.fn(async () => ({ output_text: 'AI_EXTRACTED_DESCRIPTION' })),
|
||||
getCompletionContent: vi.fn(() => 'AI_EXTRACTED_DESCRIPTION'),
|
||||
}))
|
||||
|
||||
const cosMock = vi.hoisted(() => {
|
||||
let keyIndex = 0
|
||||
return {
|
||||
generateUniqueKey: vi.fn(() => `reference-key-${++keyIndex}.jpg`),
|
||||
getSignedUrl: vi.fn((key: string) => `https://signed.example/${key}`),
|
||||
uploadObject: vi.fn(async (_buffer: Buffer, key: string) => `cos/${key}`),
|
||||
}
|
||||
})
|
||||
|
||||
const fontsMock = vi.hoisted(() => ({
|
||||
initializeFonts: vi.fn(async () => {}),
|
||||
createLabelSVG: vi.fn(async () => Buffer.from('<svg />')),
|
||||
}))
|
||||
|
||||
const workersSharedMock = vi.hoisted(() => ({
|
||||
reportTaskProgress: vi.fn(async () => {}),
|
||||
}))
|
||||
|
||||
const workersUtilsMock = vi.hoisted(() => ({
|
||||
assertTaskActive: vi.fn(async () => {}),
|
||||
}))
|
||||
|
||||
const promptI18nMock = vi.hoisted(() => ({
|
||||
PROMPT_IDS: {
|
||||
CHARACTER_IMAGE_TO_DESCRIPTION: 'character_image_to_description',
|
||||
CHARACTER_REFERENCE_TO_SHEET: 'character_reference_to_sheet',
|
||||
},
|
||||
buildPrompt: vi.fn((input: { promptId: string }) => (
|
||||
input.promptId === 'character_reference_to_sheet'
|
||||
? 'BASE_REFERENCE_PROMPT'
|
||||
: 'ANALYSIS_PROMPT'
|
||||
)),
|
||||
}))
|
||||
|
||||
const prismaMock = vi.hoisted(() => ({
|
||||
globalCharacterAppearance: {
|
||||
update: vi.fn<(input: { data?: Record<string, unknown>; where?: Record<string, unknown> }) => Promise<Record<string, never>>>(
|
||||
async () => ({}),
|
||||
),
|
||||
},
|
||||
characterAppearance: {
|
||||
update: vi.fn(async () => ({})),
|
||||
},
|
||||
}))
|
||||
|
||||
vi.mock('sharp', () => ({
|
||||
default: sharpMock,
|
||||
}))
|
||||
vi.mock('@/lib/prisma', () => ({ prisma: prismaMock }))
|
||||
vi.mock('@/lib/generator-api', () => generatorApiMock)
|
||||
vi.mock('@/lib/async-submit', () => asyncSubmitMock)
|
||||
vi.mock('@/lib/ark-api', () => arkApiMock)
|
||||
vi.mock('@/lib/api-config', () => apiConfigMock)
|
||||
vi.mock('@/lib/config-service', () => configServiceMock)
|
||||
vi.mock('@/lib/llm-client', () => llmClientMock)
|
||||
vi.mock('@/lib/storage', () => cosMock)
|
||||
vi.mock('@/lib/fonts', () => fontsMock)
|
||||
vi.mock('@/lib/workers/shared', () => workersSharedMock)
|
||||
vi.mock('@/lib/workers/utils', () => workersUtilsMock)
|
||||
vi.mock('@/lib/prompt-i18n', () => promptI18nMock)
|
||||
|
||||
import { handleReferenceToCharacterTask } from '@/lib/workers/handlers/reference-to-character'
|
||||
|
||||
function buildJob(payload: Record<string, unknown>, type: TaskType): Job<TaskJobData> {
|
||||
return {
|
||||
data: {
|
||||
taskId: 'task-1',
|
||||
type,
|
||||
locale: 'zh',
|
||||
projectId: 'project-1',
|
||||
targetType: 'GlobalCharacter',
|
||||
targetId: 'target-1',
|
||||
payload,
|
||||
userId: 'user-1',
|
||||
},
|
||||
} as unknown as Job<TaskJobData>
|
||||
}
|
||||
|
||||
function readGenerateCall(index: number) {
|
||||
const call = generatorApiMock.generateImage.mock.calls[index]
|
||||
if (!call) {
|
||||
return {
|
||||
prompt: '',
|
||||
options: {} as Record<string, unknown>,
|
||||
}
|
||||
}
|
||||
const prompt = typeof call[2] === 'string' ? call[2] : ''
|
||||
const options = (typeof call[3] === 'object' && call[3]) ? call[3] as Record<string, unknown> : {}
|
||||
return { prompt, options }
|
||||
}
|
||||
|
||||
describe('worker reference-to-character', () => {
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks()
|
||||
})
|
||||
|
||||
it('fails fast when reference images are missing', async () => {
|
||||
const job = buildJob({}, TASK_TYPE.ASSET_HUB_REFERENCE_TO_CHARACTER)
|
||||
await expect(handleReferenceToCharacterTask(job)).rejects.toThrow('Missing referenceImageUrl or referenceImageUrls')
|
||||
})
|
||||
|
||||
it('fails fast on unsupported task type', async () => {
|
||||
const job = buildJob(
|
||||
{ referenceImageUrl: 'https://example.com/ref.png' },
|
||||
'unsupported-task' as TaskType,
|
||||
)
|
||||
await expect(handleReferenceToCharacterTask(job)).rejects.toThrow('Unsupported task type')
|
||||
})
|
||||
|
||||
it('uses suffix prompt and disables reference-image injection when customDescription is provided', async () => {
|
||||
const job = buildJob(
|
||||
{
|
||||
referenceImageUrls: ['https://example.com/ref-a.png', 'https://example.com/ref-b.png'],
|
||||
customDescription: '冷静黑发角色',
|
||||
characterName: 'Hero',
|
||||
},
|
||||
TASK_TYPE.ASSET_HUB_REFERENCE_TO_CHARACTER,
|
||||
)
|
||||
|
||||
const result = await handleReferenceToCharacterTask(job)
|
||||
|
||||
expect(result).toEqual(expect.objectContaining({ success: true }))
|
||||
expect(generatorApiMock.generateImage).toHaveBeenCalledTimes(3)
|
||||
|
||||
const { prompt, options } = readGenerateCall(0)
|
||||
expect(prompt).toContain('冷静黑发角色')
|
||||
expect(prompt).toContain(CHARACTER_PROMPT_SUFFIX)
|
||||
expect(options.aspectRatio).toBe(CHARACTER_IMAGE_BANANA_RATIO)
|
||||
expect(options.referenceImages).toBeUndefined()
|
||||
})
|
||||
|
||||
it('keeps three-view suffix in template flow and writes extracted description in background mode', async () => {
|
||||
const job = buildJob(
|
||||
{
|
||||
referenceImageUrls: [' https://example.com/ref-a.png ', 'https://example.com/ref-b.png'],
|
||||
isBackgroundJob: true,
|
||||
characterId: 'character-1',
|
||||
appearanceId: 'appearance-1',
|
||||
characterName: 'Hero',
|
||||
},
|
||||
TASK_TYPE.ASSET_HUB_REFERENCE_TO_CHARACTER,
|
||||
)
|
||||
|
||||
const result = await handleReferenceToCharacterTask(job)
|
||||
|
||||
expect(result).toEqual(expect.objectContaining({ success: true }))
|
||||
expect(generatorApiMock.generateImage).toHaveBeenCalledTimes(3)
|
||||
|
||||
const { prompt, options } = readGenerateCall(0)
|
||||
expect(prompt).toContain('BASE_REFERENCE_PROMPT')
|
||||
expect(prompt).toContain(CHARACTER_PROMPT_SUFFIX)
|
||||
expect(options.referenceImages).toEqual(['https://example.com/ref-a.png', 'https://example.com/ref-b.png'])
|
||||
expect(options.aspectRatio).toBe(CHARACTER_IMAGE_BANANA_RATIO)
|
||||
|
||||
const updateArg = prismaMock.globalCharacterAppearance.update.mock.calls[0]?.[0] as {
|
||||
data?: Record<string, unknown>
|
||||
where?: Record<string, unknown>
|
||||
} | undefined
|
||||
const updateData = updateArg?.data || {}
|
||||
expect(updateArg?.where).toEqual({ id: 'appearance-1' })
|
||||
expect(updateData.description).toBe('AI_EXTRACTED_DESCRIPTION')
|
||||
expect(typeof updateData.imageUrls).toBe('string')
|
||||
expect(updateData.imageUrl).toMatch(/^cos\/reference-key-\d+\.jpg$/)
|
||||
})
|
||||
|
||||
it('uses requested count when generating reference character sheets', async () => {
|
||||
const job = buildJob(
|
||||
{
|
||||
referenceImageUrls: ['https://example.com/ref-a.png'],
|
||||
characterName: 'Hero',
|
||||
count: 5,
|
||||
},
|
||||
TASK_TYPE.REFERENCE_TO_CHARACTER,
|
||||
)
|
||||
|
||||
const result = await handleReferenceToCharacterTask(job)
|
||||
|
||||
expect(result).toEqual(expect.objectContaining({ success: true }))
|
||||
expect(generatorApiMock.generateImage).toHaveBeenCalledTimes(5)
|
||||
const cosKeys = (result as { cosKeys?: string[] }).cosKeys
|
||||
expect(cosKeys).toHaveLength(5)
|
||||
expect(cosKeys?.every((item) => item.startsWith('cos/reference-key-'))).toBe(true)
|
||||
})
|
||||
})
|
||||
Reference in New Issue
Block a user