| 1 | import { BadRequestException } from '@nestjs/common' |
| 2 | import { AppException, ResponseCode } from '@yikart/common' |
| 3 | import { AxiosError } from 'axios' |
| 4 | import { describe, expect, it, vi } from 'vitest' |
| 5 | import { isNonRetryableAiRequestError, runWithAiGenerationRetry } from './ai-generation-retry.util' |
| 6 | |
| 7 | describe('ai generation retry util', () => { |
| 8 | it('treats request errors as non-retryable', () => { |
| 9 | expect(isNonRetryableAiRequestError(new AppException(ResponseCode.InvalidModel))).toBe(true) |
| 10 | expect(isNonRetryableAiRequestError(new BadRequestException('image size is not supported'))).toBe(true) |
| 11 | }) |
| 12 | |
| 13 | it('keeps content safety errors retryable', () => { |
| 14 | expect(isNonRetryableAiRequestError(new BadRequestException('content policy violation'))).toBe(false) |
| 15 | expect(isNonRetryableAiRequestError(new Error('内容安全拦截'))).toBe(false) |
| 16 | expect(isNonRetryableAiRequestError(new Error('风控失败'))).toBe(false) |
| 17 | }) |
| 18 | |
| 19 | it('keeps provider and transport failures retryable by default', () => { |
| 20 | const rateLimitError = new AxiosError('rate limited') |
| 21 | rateLimitError.response = { status: 429 } as never |
| 22 | |
| 23 | const upstreamError = new AxiosError('upstream failed') |
| 24 | upstreamError.response = { status: 500 } as never |
| 25 | |
| 26 | expect(isNonRetryableAiRequestError(rateLimitError)).toBe(false) |
| 27 | expect(isNonRetryableAiRequestError(upstreamError)).toBe(false) |
| 28 | expect(isNonRetryableAiRequestError(new Error('socket timeout'))).toBe(false) |
| 29 | }) |
| 30 | |
| 31 | it('uses retry as extra attempts', async () => { |
| 32 | const run = vi.fn() |
| 33 | .mockRejectedValueOnce(new Error('upstream failed')) |
| 34 | .mockResolvedValue('ok') |
| 35 | |
| 36 | await expect(runWithAiGenerationRetry(run, 1)).resolves.toBe('ok') |
| 37 | expect(run).toHaveBeenCalledTimes(2) |
| 38 | }) |
| 39 | |
| 40 | it('does not retry non-retryable errors', async () => { |
| 41 | const run = vi.fn().mockRejectedValue(new AppException(ResponseCode.InvalidModel)) |
| 42 | |
| 43 | await expect(runWithAiGenerationRetry(run, 1)).rejects.toMatchObject({ code: ResponseCode.InvalidModel }) |
| 44 | expect(run).toHaveBeenCalledTimes(1) |
| 45 | }) |
| 46 | }) |
| 47 |