mirror of https://github.com/buster-so/buster.git
85 lines
2.9 KiB
TypeScript
85 lines
2.9 KiB
TypeScript
import { describe, expect, it, vi } from 'vitest';
|
|
import { CompanyResearchError } from '../../src/deep-research/types.js';
|
|
import { pollJobStatus } from '../../src/utils/polling.js';
|
|
|
|
describe('pollJobStatus', () => {
|
|
it('should return immediately when job is completed', async () => {
|
|
const mockStatus = { status: 'completed', data: 'test-data' };
|
|
const statusChecker = vi.fn().mockResolvedValue(mockStatus);
|
|
const isCompleted = vi.fn().mockReturnValue(true);
|
|
const isFailed = vi.fn().mockReturnValue(false);
|
|
const getErrorMessage = vi.fn();
|
|
|
|
const result = await pollJobStatus(
|
|
'test-job-id',
|
|
statusChecker,
|
|
isCompleted,
|
|
isFailed,
|
|
getErrorMessage,
|
|
{
|
|
interval: 1000,
|
|
maxWaitTime: 5000,
|
|
}
|
|
);
|
|
|
|
expect(result).toEqual(mockStatus);
|
|
expect(statusChecker).toHaveBeenCalledTimes(1);
|
|
expect(statusChecker).toHaveBeenCalledWith('test-job-id');
|
|
expect(isCompleted).toHaveBeenCalledWith(mockStatus);
|
|
expect(isFailed).not.toHaveBeenCalled();
|
|
});
|
|
|
|
it('should throw error when job fails', async () => {
|
|
const mockStatus = { status: 'failed', error: 'Job failed' };
|
|
const statusChecker = vi.fn().mockResolvedValue(mockStatus);
|
|
const isCompleted = vi.fn().mockReturnValue(false);
|
|
const isFailed = vi.fn().mockReturnValue(true);
|
|
const getErrorMessage = vi.fn().mockReturnValue('Job failed');
|
|
|
|
await expect(
|
|
pollJobStatus('test-job-id', statusChecker, isCompleted, isFailed, getErrorMessage, {
|
|
interval: 100,
|
|
maxWaitTime: 1000,
|
|
})
|
|
).rejects.toThrow(CompanyResearchError);
|
|
|
|
expect(statusChecker).toHaveBeenCalledTimes(1);
|
|
expect(isFailed).toHaveBeenCalledWith(mockStatus);
|
|
expect(getErrorMessage).toHaveBeenCalledWith(mockStatus);
|
|
});
|
|
|
|
it('should timeout when maxWaitTime is exceeded', async () => {
|
|
const mockStatus = { status: 'running' };
|
|
const statusChecker = vi.fn().mockResolvedValue(mockStatus);
|
|
const isCompleted = vi.fn().mockReturnValue(false);
|
|
const isFailed = vi.fn().mockReturnValue(false);
|
|
const getErrorMessage = vi.fn();
|
|
|
|
await expect(
|
|
pollJobStatus('test-job-id', statusChecker, isCompleted, isFailed, getErrorMessage, {
|
|
interval: 100,
|
|
maxWaitTime: 300, // Very short timeout
|
|
})
|
|
).rejects.toThrow(CompanyResearchError);
|
|
|
|
// Should have been called multiple times before timeout
|
|
expect(statusChecker).toHaveBeenCalledWith('test-job-id');
|
|
});
|
|
|
|
it('should handle statusChecker errors', async () => {
|
|
const statusChecker = vi.fn().mockRejectedValue(new Error('Network error'));
|
|
const isCompleted = vi.fn();
|
|
const isFailed = vi.fn();
|
|
const getErrorMessage = vi.fn();
|
|
|
|
await expect(
|
|
pollJobStatus('test-job-id', statusChecker, isCompleted, isFailed, getErrorMessage, {
|
|
interval: 100,
|
|
maxWaitTime: 1000,
|
|
})
|
|
).rejects.toThrow(CompanyResearchError);
|
|
|
|
expect(statusChecker).toHaveBeenCalledTimes(1);
|
|
});
|
|
});
|