ai-functions
Version:
Core AI primitives for building intelligent applications
353 lines (352 loc) • 15.4 kB
JavaScript
/**
* Tests for batch and background processing modes
*
* Batch mode: fn.batch([inputs]) - Process many inputs at ~50% discount
* Background mode: fn(..., { mode: 'background' }) - Returns job ID immediately
*/
import { describe, it, expect, vi, beforeEach } from 'vitest';
// ============================================================================
// Mock implementations
// ============================================================================
const mockBatchProcess = vi.fn();
const mockBackgroundProcess = vi.fn();
/**
* Create a mock function with batch support
*/
function createMockFunctionWithBatch(defaultHandler) {
const fn = async (prompt, options) => {
if (options?.mode === 'background') {
return mockBackgroundProcess(prompt, options);
}
return defaultHandler(prompt);
};
// Add batch method
fn.batch = async (inputs) => {
return mockBatchProcess(inputs);
};
return fn;
}
/**
* Create a tagged template function with batch support
*/
function createMockTemplateFunctionWithBatch(defaultHandler) {
function fn(promptOrStrings, ...args) {
let prompt;
if (Array.isArray(promptOrStrings) && 'raw' in promptOrStrings) {
prompt = promptOrStrings.reduce((acc, str, i) => {
return acc + str + (args[i] ?? '');
}, '');
}
else {
prompt = promptOrStrings;
}
return defaultHandler(prompt);
}
// Add batch method
fn.batch = async (inputs) => {
return mockBatchProcess(inputs);
};
return fn;
}
// ============================================================================
// Batch mode tests
// ============================================================================
describe('batch mode', () => {
beforeEach(() => {
mockBatchProcess.mockReset();
});
describe('write.batch()', () => {
it('processes multiple prompts in batch', async () => {
const write = createMockFunctionWithBatch(async () => 'Generated content');
const prompts = [
'blog post about TypeScript',
'blog post about React',
'blog post about Next.js',
];
mockBatchProcess.mockResolvedValue([
'TypeScript content...',
'React content...',
'Next.js content...',
]);
const results = await write.batch(prompts);
expect(mockBatchProcess).toHaveBeenCalledWith(prompts);
expect(results).toHaveLength(3);
});
it('processes object inputs with context', async () => {
const write = createMockFunctionWithBatch(async () => 'Generated content');
const brand = { voice: 'professional', audience: 'developers' };
const titles = ['Getting Started', 'Advanced Patterns', 'Best Practices'];
const inputs = titles.map(title => ({
title,
brand,
tone: 'technical',
}));
mockBatchProcess.mockResolvedValue([
'Getting Started content...',
'Advanced Patterns content...',
'Best Practices content...',
]);
const results = await write.batch(inputs);
expect(mockBatchProcess).toHaveBeenCalledWith(inputs);
expect(results).toHaveLength(3);
});
it('returns results in same order as inputs', async () => {
const write = createMockFunctionWithBatch(async () => 'content');
const inputs = ['first', 'second', 'third'];
mockBatchProcess.mockResolvedValue([
'Result for first',
'Result for second',
'Result for third',
]);
const results = await write.batch(inputs);
expect(results[0]).toContain('first');
expect(results[1]).toContain('second');
expect(results[2]).toContain('third');
});
});
describe('list.batch()', () => {
it('generates multiple lists in batch', async () => {
const list = createMockFunctionWithBatch(async () => ['item']);
mockBatchProcess.mockResolvedValue([
['TypeScript tip 1', 'TypeScript tip 2'],
['React tip 1', 'React tip 2'],
['Next.js tip 1', 'Next.js tip 2'],
]);
const results = await list.batch([
'3 TypeScript tips',
'3 React tips',
'3 Next.js tips',
]);
expect(results).toHaveLength(3);
expect(results[0]).toEqual(['TypeScript tip 1', 'TypeScript tip 2']);
});
});
describe('code.batch()', () => {
it('generates multiple code snippets in batch', async () => {
const code = createMockFunctionWithBatch(async () => 'code');
mockBatchProcess.mockResolvedValue([
'function validateEmail(email) { ... }',
'function validatePhone(phone) { ... }',
'function validateUrl(url) { ... }',
]);
const results = await code.batch([
{ description: 'email validator', language: 'typescript' },
{ description: 'phone validator', language: 'typescript' },
{ description: 'url validator', language: 'typescript' },
]);
expect(results).toHaveLength(3);
expect(results[0]).toContain('validateEmail');
});
});
describe('batch with options', () => {
it('accepts batch-level options', async () => {
const write = createMockFunctionWithBatch(async () => 'content');
// Simulating batch with options
const mockBatchWithOptions = vi.fn().mockResolvedValue(['r1', 'r2']);
const inputs = ['prompt1', 'prompt2'];
const options = { model: 'claude-opus-4-5' };
await mockBatchWithOptions(inputs, options);
expect(mockBatchWithOptions).toHaveBeenCalledWith(inputs, options);
});
it('supports priority option for urgent batches', async () => {
const mockBatchWithPriority = vi.fn().mockResolvedValue(['result']);
await mockBatchWithPriority(['prompt'], { priority: 'high' });
expect(mockBatchWithPriority).toHaveBeenCalledWith(['prompt'], expect.objectContaining({ priority: 'high' }));
});
});
});
// ============================================================================
// Background mode tests
// ============================================================================
describe('background mode', () => {
beforeEach(() => {
mockBackgroundProcess.mockReset();
});
it('returns job ID immediately', async () => {
const write = createMockFunctionWithBatch(async () => 'content');
mockBackgroundProcess.mockResolvedValue({
jobId: 'job_abc123',
status: 'pending',
});
const job = await write('long form article', { mode: 'background' });
expect(mockBackgroundProcess).toHaveBeenCalledWith('long form article', expect.objectContaining({ mode: 'background' }));
expect(job).toHaveProperty('jobId');
expect(job.status).toBe('pending');
});
it('can check job status', async () => {
const mockGetJobStatus = vi.fn();
// Simulating job status check
mockGetJobStatus.mockResolvedValueOnce({ status: 'processing' });
mockGetJobStatus.mockResolvedValueOnce({ status: 'completed', result: 'Generated content' });
const status1 = await mockGetJobStatus('job_abc123');
expect(status1.status).toBe('processing');
const status2 = await mockGetJobStatus('job_abc123');
expect(status2.status).toBe('completed');
expect(status2.result).toBe('Generated content');
});
it('supports webhook callback', async () => {
const write = createMockFunctionWithBatch(async () => 'content');
mockBackgroundProcess.mockResolvedValue({
jobId: 'job_xyz789',
status: 'pending',
});
await write('content', {
mode: 'background',
webhook: 'https://myapp.com/webhooks/ai-complete',
});
expect(mockBackgroundProcess).toHaveBeenCalledWith('content', expect.objectContaining({
mode: 'background',
webhook: 'https://myapp.com/webhooks/ai-complete',
}));
});
it('supports polling for result', async () => {
// Simulating a poll function
const mockPollForResult = vi.fn();
mockPollForResult.mockImplementation(async (jobId) => {
// Simulate polling - would normally check periodically
return { status: 'completed', result: 'Final result' };
});
const result = await mockPollForResult('job_abc123');
expect(result.status).toBe('completed');
expect(result.result).toBe('Final result');
});
});
// ============================================================================
// Combined batch and background
// ============================================================================
describe('batch + background mode', () => {
it('can run batch in background', async () => {
const mockBatchBackground = vi.fn();
mockBatchBackground.mockResolvedValue({
jobId: 'batch_job_123',
status: 'pending',
inputCount: 100,
});
// Large batch job in background
const job = await mockBatchBackground(Array(100).fill('Generate content'), { mode: 'background' });
expect(job.jobId).toBe('batch_job_123');
expect(job.inputCount).toBe(100);
});
it('tracks progress of background batch', async () => {
const mockBatchProgress = vi.fn();
mockBatchProgress
.mockResolvedValueOnce({ status: 'processing', completed: 10, total: 100 })
.mockResolvedValueOnce({ status: 'processing', completed: 50, total: 100 })
.mockResolvedValueOnce({ status: 'completed', completed: 100, total: 100 });
const p1 = await mockBatchProgress('batch_job_123');
expect(p1.completed).toBe(10);
const p2 = await mockBatchProgress('batch_job_123');
expect(p2.completed).toBe(50);
const p3 = await mockBatchProgress('batch_job_123');
expect(p3.status).toBe('completed');
});
});
// ============================================================================
// Batch pricing and limits
// ============================================================================
describe('batch characteristics', () => {
it('batch provides cost savings (documentation test)', () => {
// This documents expected behavior
const batchInfo = {
discount: '50%',
turnaround: '24 hours max',
minBatchSize: 1,
maxBatchSize: 10000,
};
expect(batchInfo.discount).toBe('50%');
expect(batchInfo.turnaround).toBe('24 hours max');
});
it('batch handles errors gracefully', async () => {
const mockBatchWithErrors = vi.fn();
// Some items succeed, some fail
mockBatchWithErrors.mockResolvedValue({
results: [
{ index: 0, status: 'success', result: 'Content 1' },
{ index: 1, status: 'error', error: 'Content policy violation' },
{ index: 2, status: 'success', result: 'Content 3' },
],
summary: { succeeded: 2, failed: 1 },
});
const response = await mockBatchWithErrors(['p1', 'p2', 'p3']);
expect(response.summary.succeeded).toBe(2);
expect(response.summary.failed).toBe(1);
expect(response.results[1].status).toBe('error');
});
});
// ============================================================================
// Use cases from README
// ============================================================================
describe('batch use cases', () => {
beforeEach(() => {
mockBatchProcess.mockReset();
});
it('content generation at scale', async () => {
const write = createMockFunctionWithBatch(async () => 'content');
// Generate blog posts for 100 topics
const topics = Array(100)
.fill(null)
.map((_, i) => `Topic ${i + 1}`);
const inputs = topics.map(topic => ({
topic,
length: 'medium',
style: 'informative',
}));
mockBatchProcess.mockResolvedValue(topics.map(t => `Blog post about ${t}...`));
const posts = await write.batch(inputs);
expect(posts).toHaveLength(100);
});
it('product description generation', async () => {
const write = createMockFunctionWithBatch(async () => 'description');
const products = [
{ name: 'Widget Pro', category: 'tools', features: ['durable', 'lightweight'] },
{ name: 'Gadget Plus', category: 'electronics', features: ['wireless', 'rechargeable'] },
];
mockBatchProcess.mockResolvedValue([
'Widget Pro is a durable, lightweight tool...',
'Gadget Plus is a wireless, rechargeable electronic...',
]);
const descriptions = await write.batch(products.map(p => ({
prompt: `product description for ${p.name}`,
product: p,
})));
expect(descriptions).toHaveLength(2);
expect(descriptions[0]).toContain('Widget Pro');
});
it('code generation for multiple functions', async () => {
const code = createMockFunctionWithBatch(async () => 'code');
const functions = [
{ name: 'validateEmail', description: 'Validate email format' },
{ name: 'validatePhone', description: 'Validate phone number' },
{ name: 'validateUrl', description: 'Validate URL format' },
];
mockBatchProcess.mockResolvedValue(functions.map(f => `function ${f.name}(value) { ... }`));
const implementations = await code.batch(functions.map(f => ({
description: f.description,
functionName: f.name,
language: 'typescript',
})));
expect(implementations).toHaveLength(3);
functions.forEach((f, i) => {
expect(implementations[i]).toContain(f.name);
});
});
});
// ============================================================================
// Error handling
// ============================================================================
describe('batch error handling', () => {
beforeEach(() => {
mockBatchProcess.mockReset();
});
it('handles empty input array', async () => {
const write = createMockFunctionWithBatch(async () => 'content');
mockBatchProcess.mockResolvedValue([]);
const results = await write.batch([]);
expect(results).toEqual([]);
});
it('propagates batch-level errors', async () => {
const write = createMockFunctionWithBatch(async () => 'content');
mockBatchProcess.mockRejectedValue(new Error('Batch quota exceeded'));
await expect(write.batch(['prompt'])).rejects.toThrow('Batch quota exceeded');
});
});