| 1 | import type { AiImageData, DraftGenerationData } from './interfaces' |
| 2 | import { UserType } from '@yikart/common' |
| 3 | import { describe, expect, it, vi } from 'vitest' |
| 4 | import { QueueService } from './queue.service' |
| 5 | |
| 6 | describe('queueService draft generation queue info', () => { |
| 7 | function createService(queues?: { |
| 8 | aiImageAsync?: Record<string, unknown> |
| 9 | normal?: Record<string, unknown> |
| 10 | lowPriority?: Record<string, unknown> |
| 11 | }) { |
| 12 | const aiImageAsyncQueue = { |
| 13 | add: vi.fn(), |
| 14 | getJob: vi.fn(), |
| 15 | getJobs: vi.fn(), |
| 16 | ...queues?.aiImageAsync, |
| 17 | } |
| 18 | const normalQueue = { |
| 19 | add: vi.fn(), |
| 20 | getJob: vi.fn(), |
| 21 | count: vi.fn(), |
| 22 | getJobs: vi.fn(), |
| 23 | ...queues?.normal, |
| 24 | } |
| 25 | const lowPriorityQueue = { |
| 26 | add: vi.fn(), |
| 27 | getJob: vi.fn(), |
| 28 | count: vi.fn(), |
| 29 | getJobs: vi.fn(), |
| 30 | ...queues?.lowPriority, |
| 31 | } |
| 32 | const service = Object.create(QueueService.prototype) as { |
| 33 | defaultOptions: Record<string, unknown> |
| 34 | aiImageAsyncQueue: typeof aiImageAsyncQueue |
| 35 | draftGenerationQueue: typeof normalQueue |
| 36 | draftGenerationLowPriorityQueue: typeof lowPriorityQueue |
| 37 | addAiImageAsyncJob: QueueService['addAiImageAsyncJob'] |
| 38 | isAiImageAsyncJobActive: QueueService['isAiImageAsyncJobActive'] |
| 39 | addDraftGenerationJob: QueueService['addDraftGenerationJob'] |
| 40 | addLowPriorityDraftGenerationJob: QueueService['addLowPriorityDraftGenerationJob'] |
| 41 | getDraftGenerationQueueInfo: QueueService['getDraftGenerationQueueInfo'] |
| 42 | isDraftGenerationJobActive: QueueService['isDraftGenerationJobActive'] |
| 43 | } |
| 44 | |
| 45 | service.defaultOptions = { removeOnComplete: { age: 30, count: 1000 } } |
| 46 | service.aiImageAsyncQueue = aiImageAsyncQueue |
| 47 | service.draftGenerationQueue = normalQueue |
| 48 | service.draftGenerationLowPriorityQueue = lowPriorityQueue |
| 49 | |
| 50 | return { service, aiImageAsyncQueue, normalQueue, lowPriorityQueue } |
| 51 | } |
| 52 | |
| 53 | const data: DraftGenerationData = { |
| 54 | aiLogId: 'log-1', |
| 55 | userId: 'user-1', |
| 56 | userType: UserType.User, |
| 57 | groupId: 'group-1', |
| 58 | version: 'v2-image-text', |
| 59 | } |
| 60 | |
| 61 | const imageData = { |
| 62 | logId: 'image-log-1', |
| 63 | userId: 'user-1', |
| 64 | userType: UserType.User, |
| 65 | model: 'seedream-3', |
| 66 | type: 'image', |
| 67 | request: {}, |
| 68 | taskType: 'generation', |
| 69 | } as unknown as AiImageData |
| 70 | |
| 71 | it('异步图片入队时固定使用 logId 作为 jobId', async () => { |
| 72 | const { service, aiImageAsyncQueue } = createService() |
| 73 | |
| 74 | await service.addAiImageAsyncJob(imageData) |
| 75 | |
| 76 | expect(aiImageAsyncQueue.add).toHaveBeenCalledWith('generate', imageData, expect.objectContaining({ |
| 77 | jobId: 'image-log-1', |
| 78 | })) |
| 79 | }) |
| 80 | |
| 81 | it('异步图片 active 查询优先使用固定 jobId', async () => { |
| 82 | const getState = vi.fn().mockResolvedValue('active') |
| 83 | const { service, aiImageAsyncQueue } = createService({ |
| 84 | aiImageAsync: { |
| 85 | getJob: vi.fn().mockResolvedValue({ getState }), |
| 86 | getJobs: vi.fn(), |
| 87 | }, |
| 88 | }) |
| 89 | |
| 90 | await expect(service.isAiImageAsyncJobActive('image-log-1')).resolves.toBe(true) |
| 91 | |
| 92 | expect(aiImageAsyncQueue.getJob).toHaveBeenCalledWith('image-log-1') |
| 93 | expect(aiImageAsyncQueue.getJobs).not.toHaveBeenCalled() |
| 94 | }) |
| 95 | |
| 96 | it('异步图片 active 查询兼容旧任务扫描 logId', async () => { |
| 97 | const { service, aiImageAsyncQueue } = createService({ |
| 98 | aiImageAsync: { |
| 99 | getJob: vi.fn().mockResolvedValue(undefined), |
| 100 | getJobs: vi.fn().mockResolvedValue([ |
| 101 | { data: { logId: 'other-log' } }, |
| 102 | { data: { logId: 'image-log-1' } }, |
| 103 | ]), |
| 104 | }, |
| 105 | }) |
| 106 | |
| 107 | await expect(service.isAiImageAsyncJobActive('image-log-1')).resolves.toBe(true) |
| 108 | |
| 109 | expect(aiImageAsyncQueue.getJobs).toHaveBeenCalledWith(['active'], 0, -1, true) |
| 110 | }) |
| 111 | |
| 112 | it('草稿生成入队时固定使用 aiLogId 作为 jobId', async () => { |
| 113 | const { service, normalQueue, lowPriorityQueue } = createService() |
| 114 | |
| 115 | await service.addDraftGenerationJob(data, { priority: 1000 }) |
| 116 | await service.addLowPriorityDraftGenerationJob(data, { priority: 1000 }) |
| 117 | |
| 118 | expect(normalQueue.add).toHaveBeenCalledWith('generate', data, expect.objectContaining({ |
| 119 | priority: 1000, |
| 120 | jobId: 'log-1', |
| 121 | attempts: 3, |
| 122 | backoff: { |
| 123 | type: 'exponential', |
| 124 | delay: 5000, |
| 125 | }, |
| 126 | })) |
| 127 | expect(lowPriorityQueue.add).toHaveBeenCalledWith('generate', data, expect.objectContaining({ |
| 128 | priority: 1000, |
| 129 | jobId: 'log-1', |
| 130 | attempts: 3, |
| 131 | backoff: { |
| 132 | type: 'exponential', |
| 133 | delay: 5000, |
| 134 | }, |
| 135 | })) |
| 136 | }) |
| 137 | |
| 138 | it('返回队列位置时只包含 position 和 waitingCount', async () => { |
| 139 | const { service, normalQueue } = createService({ |
| 140 | normal: { |
| 141 | getJob: vi.fn().mockResolvedValue({ |
| 142 | id: 'log-1', |
| 143 | getState: vi.fn().mockResolvedValue('prioritized'), |
| 144 | }), |
| 145 | count: vi.fn().mockResolvedValue(8), |
| 146 | getJobs: vi.fn().mockResolvedValue([ |
| 147 | { id: 'log-0' }, |
| 148 | { id: 'log-1' }, |
| 149 | ]), |
| 150 | }, |
| 151 | }) |
| 152 | |
| 153 | const queueInfo = await service.getDraftGenerationQueueInfo('log-1') |
| 154 | |
| 155 | expect(queueInfo).toEqual({ |
| 156 | position: 2, |
| 157 | waitingCount: 8, |
| 158 | }) |
| 159 | expect(Object.keys(queueInfo ?? {})).toEqual(['position', 'waitingCount']) |
| 160 | expect(normalQueue.getJobs).toHaveBeenCalledWith(['prioritized', 'waiting', 'delayed', 'waiting-children'], 0, -1, true) |
| 161 | }) |
| 162 | |
| 163 | it('普通队列未命中时查询低优先级队列但不暴露队列类型', async () => { |
| 164 | const { service, lowPriorityQueue } = createService({ |
| 165 | normal: { |
| 166 | getJob: vi.fn().mockResolvedValue(undefined), |
| 167 | }, |
| 168 | lowPriority: { |
| 169 | getJob: vi.fn().mockResolvedValue({ |
| 170 | id: 'log-1', |
| 171 | getState: vi.fn().mockResolvedValue('active'), |
| 172 | }), |
| 173 | count: vi.fn().mockResolvedValue(4), |
| 174 | getJobs: vi.fn(), |
| 175 | }, |
| 176 | }) |
| 177 | |
| 178 | const queueInfo = await service.getDraftGenerationQueueInfo('log-1') |
| 179 | |
| 180 | expect(queueInfo).toEqual({ |
| 181 | position: null, |
| 182 | waitingCount: 4, |
| 183 | }) |
| 184 | expect(lowPriorityQueue.getJobs).not.toHaveBeenCalled() |
| 185 | }) |
| 186 | |
| 187 | it('草稿 active 查询检查普通和低优先级队列', async () => { |
| 188 | const { service, normalQueue, lowPriorityQueue } = createService({ |
| 189 | normal: { |
| 190 | getJob: vi.fn().mockResolvedValue({ |
| 191 | getState: vi.fn().mockResolvedValue('completed'), |
| 192 | }), |
| 193 | }, |
| 194 | lowPriority: { |
| 195 | getJob: vi.fn().mockResolvedValue({ |
| 196 | getState: vi.fn().mockResolvedValue('active'), |
| 197 | }), |
| 198 | }, |
| 199 | }) |
| 200 | |
| 201 | await expect(service.isDraftGenerationJobActive('log-1')).resolves.toBe(true) |
| 202 | |
| 203 | expect(normalQueue.getJob).toHaveBeenCalledWith('log-1') |
| 204 | expect(lowPriorityQueue.getJob).toHaveBeenCalledWith('log-1') |
| 205 | }) |
| 206 | }) |
| 207 |