Unit Testing AI Components: Patterns That Don't Make You Want to Quit

The first time you try to unit test AI components, you hit a wall: LLM behavior is non-deterministic. You can't assert "the output is exactly this." You can't reliably mock responses. Your tests pass locally but fail in CI. And if you're calling the real API, tests are slow and expensive.

The Architecture Shift: Separating AI from Business Logic

The first thing to understand: AI logic and business logic are different beasts. Test them separately.

Bad approach: writing tests that call the real LLM and assert on the output.

Good approach: splitting your code so you test business logic in isolation, then test AI integration points separately with mocks.

Here's what this looks like:

// BAD: AI logic and business logic mixed
async function generateRecommendationBad(userId) {
 const user = await db.getUser(userId);
 const history = await db.getUserHistory(userId);

 // Call real API
 const response = await client.messages.create({
 model: "claude-3-5-sonnet-20241022",
 messages: [
 {
 role: "user",
 content: `User ${user.name} has viewed ${history.length} items...`
 }
 ]
 });

 const recommendation = parseRecommendation(response);
 await db.saveRecommendation(userId, recommendation);
 return recommendation;
}

// GOOD: Separated concerns
class RecommendationService {
 constructor(private llmProvider: LLMProvider, private db: Database) {}

 async generateRecommendation(userId: string) {
 // Step 1: Load data (testable, no AI)
 const user = await this.db.getUser(userId);
 const history = await this.db.getUserHistory(userId);
 const context = this.buildContext(user, history);

 // Step 2: Call AI (mockable, testable separately)
 const aiOutput = await this.getAIRecommendation(context);

 // Step 3: Process output (pure logic, fully testable)
 const recommendation = this.parseRecommendation(aiOutput);

 // Step 4: Persist (testable with mocks)
 await this.db.saveRecommendation(userId, recommendation);

 return recommendation;
 }

 // Isolate AI call so it can be mocked
 private async getAIRecommendation(context: UserContext) {
 return this.llmProvider.generateText({
 prompt: this.buildPrompt(context),
 maxTokens: 500
 });
 }

 // Test this without AI
 private buildContext(user: User, history: UserHistory[]): UserContext {
 return {
 name: user.name,
 interests: history.map(h => h.category),
 purchaseFrequency: this.calculateFrequency(history)
 };
 }

 // Test this without AI
 private parseRecommendation(output: string): Recommendation {
 // Pure logic to extract recommendation from text
 const match = output.match(/recommendation:\s*(.+)/i);
 return {
 id: match?.[1] || "",
 confidence: this.calculateConfidence(output)
 };
 }
}

Now you can test:

• buildContext(), pure function, no mocks needed
• parseRecommendation(), pure function, test with fixtures
• getAIRecommendation(), mock the LLM provider, test integration
• generateRecommendation(), mock database and LLM, test flow

Mocking LLM Responses: The Right Way

Never mock random responses. Never call the real API in tests. Use fixture-based mocking with realistic, reproducible responses.

// Create an LLM provider interface
interface LLMProvider {
 generateText(request: GenerateRequest): Promise;
}

// Real implementation calls actual API
class AnthropicLLMProvider implements LLMProvider {
 async generateText(request: GenerateRequest): Promise {
 const response = await this.client.messages.create({
 model: "claude-3-5-sonnet-20241022",
 max_tokens: request.maxTokens,
 messages: [{ role: "user", content: request.prompt }]
 });

 return response.content[0].text;
 }
}

// Mock implementation returns fixtures
class MockLLMProvider implements LLMProvider {
 private fixtures: Map = new Map();

 addFixture(prompt: string, response: string) {
 this.fixtures.set(this.normalizePrompt(prompt), response);
 }

 async generateText(request: GenerateRequest): Promise {
 const normalized = this.normalizePrompt(request.prompt);
 const fixture = this.fixtures.get(normalized);

 if (!fixture) {
 throw new Error(
 `No fixture for prompt: ${request.prompt}. ` +
 `Add it to your mock provider.`
 );
 }

 return fixture;
 }

 private normalizePrompt(prompt: string): string {
 // Normalize to handle minor variations
 return prompt
 .toLowerCase()
 .replace(/\s+/g, " ")
 .trim();
 }
}

// Use in tests
describe("RecommendationService", () => {
 let service: RecommendationService;
 let mockLLM: MockLLMProvider;
 let mockDB: jest.Mocked;

 beforeEach(() => {
 mockLLM = new MockLLMProvider();
 mockDB = jest.mocked(new Database());

 service = new RecommendationService(mockLLM, mockDB);
 });

 it("should generate recommendation from user history", async () => {
 // Set up fixtures
 mockLLM.addFixture(
 "User alice has viewed science fiction, technology, business. What should we recommend?",
 "recommendation: dune"
 );

 mockDB.getUser.mockResolvedValue({
 id: "user_1",
 name: "alice"
 });

 mockDB.getUserHistory.mockResolvedValue([
 { category: "science fiction" },
 { category: "technology" },
 { category: "business" }
 ]);

 const result = await service.generateRecommendation("user_1");

 expect(result.id).toBe("dune");
 expect(mockDB.saveRecommendation).toHaveBeenCalledWith(
 "user_1",
 expect.objectContaining({ id: "dune" })
 );
 });

 it("should handle missing fixture gracefully", async () => {
 mockDB.getUser.mockResolvedValue({
 id: "user_2",
 name: "bob"
 });

 mockDB.getUserHistory.mockResolvedValue([
 { category: "unknown_category" }
 ]);

 // Should throw because fixture doesn't exist
 await expect(
 service.generateRecommendation("user_2")
 ).rejects.toThrow("No fixture for prompt");
 });
});

Key benefits of fixture-based mocking:

• Deterministic: Same input always produces same output
• Fast: No API calls, tests run instantly
• Cheap: No token usage, free to run thousands of times
• Realistic: You're testing with actual LLM outputs, not random data
• Catches integration gaps: If the real API changes, fixtures force you to update

Testing Prompt Templates

Your prompts are contracts. They define what you're asking the model to do. Test them as data structures, not as strings.

// Define prompts as structured templates
class PromptTemplate {
 constructor(
 private template: string,
 private variables: string[],
 private expectedOutputFormat: "json" | "text" | "markdown"
 ) {}

 render(values: Record): string {
 let result = this.template;

 for (const variable of this.variables) {
 const regex = new RegExp(`{{${variable}}}`, "g");
 const value = values[variable];

 if (!value) {
 throw new Error(`Missing required variable: ${variable}`);
 }

 result = result.replace(regex, value);
 }

 return result;
 }

 // Test that prompt renders correctly
 validate(values: Record): boolean {
 try {
 this.render(values);
 return true;
 } catch {
 return false;
 }
 }
}

// Define your prompts
const classificationPrompt = new PromptTemplate(
 `You are a content moderator. Classify the following user message as safe or unsafe.

 User message: {{message}}

 Respond with ONLY "safe" or "unsafe", followed by a confidence score 0-1.`,
 ["message"],
 "text"
);

const summarizationPrompt = new PromptTemplate(
 `Summarize the following text in 2-3 sentences. Focus on key insights.

 Text: {{text}}

 Summary:`,
 ["text"],
 "text"
);

// Test prompts as data
describe("Prompts", () => {
 it("should have all required variables", () => {
 const rendered = classificationPrompt.render({
 message: "This is a test message"
 });

 expect(rendered).toContain("test message");
 expect(rendered).not.toContain("{{");
 });

 it("should throw on missing variables", () => {
 expect(() => {
 classificationPrompt.render({});
 }).toThrow("Missing required variable: message");
 });

 it("should handle special characters in variables", () => {
 const dangerous =
 'This message has "quotes" and \\n newlines and {curly} braces';

 const rendered = classificationPrompt.render({
 message: dangerous
 });

 expect(rendered).toContain(dangerous);
 });

 it("should not have template syntax in rendered output", () => {
 const rendered = classificationPrompt.render({
 message: "test"
 });

 expect(rendered).not.toMatch(/{{.*}}/);
 });
});

// Use in service
class ModerationService {
 async classifyMessage(message: string): Promise {
 const prompt = classificationPrompt.render({ message });

 const response = await this.llmProvider.generateText({
 prompt,
 maxTokens: 50
 });

 return this.parseClassification(response);
 }

 private parseClassification(output: string): Classification {
 const match = output.match(/(safe|unsafe)\s+([\d.]+)/i);
 return {
 status: match?.[1].toLowerCase() as "safe" | "unsafe",
 confidence: parseFloat(match?.[2] || "0")
 };
 }
}

Dependency Injection for Model Swapping

Test against different models without changing your application code. Use dependency injection.

"Your service shouldn't know or care which LLM it's using. Inject it at construction time. This makes testing and switching models trivial."

// Interface defines what any LLM provider must support
interface ModelProvider {
 generateText(prompt: string, maxTokens: number): Promise;
 getModel(): string;
}

// Implementations for different models
class ClaudeSonnetProvider implements ModelProvider {
 getModel() {
 return "claude-3-5-sonnet-20241022";
 }

 async generateText(
 prompt: string,
 maxTokens: number
 ): Promise {
 const response = await this.client.messages.create({
 model: this.getModel(),
 max_tokens: maxTokens,
 messages: [{ role: "user", content: prompt }]
 });

 return response.content[0].text;
 }
}

class GPT4Provider implements ModelProvider {
 getModel() {
 return "gpt-4-turbo";
 }

 async generateText(
 prompt: string,
 maxTokens: number
 ): Promise {
 const response = await this.openai.chat.completions.create({
 model: this.getModel(),
 messages: [{ role: "user", content: prompt }],
 max_tokens: maxTokens
 });

 return response.choices[0].message.content || "";
 }
}

class TestProvider implements ModelProvider {
 private fixtures: Map = new Map();

 getModel() {
 return "mock-model";
 }

 addFixture(prompt: string, response: string) {
 this.fixtures.set(prompt, response);
 }

 async generateText(
 prompt: string,
 maxTokens: number
 ): Promise {
 const fixture = this.fixtures.get(prompt);
 if (!fixture) {
 throw new Error(`No fixture for: ${prompt}`);
 }
 return fixture;
 }
}

// Service doesn't know which provider it gets
class AnalysisService {
 constructor(private modelProvider: ModelProvider) {}

 async analyze(text: string): Promise {
 const response = await this.modelProvider.generateText(
 `Analyze this text: ${text}`,
 500
 );

 return this.parseAnalysis(response);
 }
}

// Use different providers at different times
describe("AnalysisService", () => {
 it("should work with any provider", async () => {
 const testProvider = new TestProvider();
 testProvider.addFixture(
 "Analyze this text: hello world",
 "sentiment: positive"
 );

 const service = new AnalysisService(testProvider);
 const result = await service.analyze("hello world");

 expect(result).toBeDefined();
 });
});

// In production
const production = new AnalysisService(new ClaudeSonnetProvider());

// For testing different models
const testWithGPT = new AnalysisService(new GPT4Provider());
const testWithClaude = new AnalysisService(new ClaudeSonnetProvider());
const testWithMock = new AnalysisService(new TestProvider());

Building Test Fixtures from Real Outputs

Start your test fixtures with real LLM outputs. Run your code against the real API, capture the response, and add it to your fixtures.

// Tool to capture and save fixtures
async function captureFixture(
 prompt: string,
 fixtureFile: string
) {
 console.log("Capturing fixture for prompt:", prompt);

 const response = await client.messages.create({
 model: "claude-3-5-sonnet-20241022",
 max_tokens: 500,
 messages: [{ role: "user", content: prompt }]
 });

 const output = response.content[0].text;

 // Save to fixtures file
 const fixtures = JSON.parse(
 fs.readFileSync(fixtureFile, "utf-8")
 );

 fixtures[prompt] = output;

 fs.writeFileSync(fixtureFile, JSON.stringify(fixtures, null, 2));

 console.log("Fixture saved");
 return output;
}

// Workflow: develop feature with real API, then commit fixtures
async function developWithRealtimeAPI() {
 // 1. Write your feature code
 // 2. Test against real API
 const realResponse = await analyzeWithRealAPI(testData);

 // 3. When satisfied, capture fixture
 await captureFixture(
 "Analyze: " + testData,
 "src/__fixtures__/analysis-fixtures.json"
 );

 // 4. Update your test to use fixture
 // 5. Commit fixture file to git
 // 6. All future tests use the fixture
}

Testing Error Cases and Edge Cases

Your tests should handle:

• Empty/null inputs
• Malformed LLM responses
• Timeout/rate limit errors
• Parsing failures

describe("RecommendationService error handling", () => {
 let service: RecommendationService;
 let mockLLM: MockLLMProvider;

 beforeEach(() => {
 mockLLM = new MockLLMProvider();
 service = new RecommendationService(mockLLM, mockDB);
 });

 it("should handle empty user history", async () => {
 mockDB.getUserHistory.mockResolvedValue([]);

 // Service should either return default or handle gracefully
 const result = await service.generateRecommendation("user_empty");

 expect(result).toBeDefined();
 });

 it("should handle malformed LLM response", async () => {
 // LLM returns something that can't be parsed
 mockLLM.addFixture("some prompt", "completely invalid response");

 mockDB.getUser.mockResolvedValue({
 id: "user_1",
 name: "test"
 });
 mockDB.getUserHistory.mockResolvedValue([
 { category: "test" }
 ]);

 // Should throw or return safe default
 await expect(
 service.generateRecommendation("user_1")
 ).rejects.toThrow();
 });

 it("should handle LLM timeout", async () => {
 const timedOutProvider = {
 generateText: jest.fn().mockRejectedValue(
 new Error("LLM timeout after 30s")
 )
 };

 const timedOutService = new RecommendationService(
 timedOutProvider,
 mockDB
 );

 await expect(
 timedOutService.generateRecommendation("user_1")
 ).rejects.toThrow("timeout");
 });
});

Integration Tests: Where Real APIs Matter

Unit tests use mocks. But you still need integration tests that call the real API. Just do it separately.

// __tests__/integration/recommendations.integration.test.ts
// Marked to run separately, not in CI pre-commit

describe(
 "RecommendationService Integration",
 () => {
 let service: RecommendationService;

 beforeEach(() => {
 // Use real providers
 const realLLM = new AnthropicLLMProvider();
 const realDB = new Database();

 service = new RecommendationService(realLLM, realDB);
 });

 it(
 "should generate valid recommendations against real API",
 async () => {
 const userId = "test_user_" + Date.now();

 await realDB.createTestUser(userId);
 await realDB.addViewToHistory(userId, "science fiction");
 await realDB.addViewToHistory(userId, "philosophy");

 const recommendation = await service.generateRecommendation(
 userId
 );

 // Test against real API output
 expect(recommendation.id).toBeTruthy();
 expect(recommendation.confidence).toBeGreaterThan(0);

 // Clean up
 await realDB.deleteTestUser(userId);
 },
 30000
 ); // Allow 30 seconds for API call
 },
 { timeout: 60000 }
);

Run integration tests:

• Manually during development
• On a schedule (nightly, not per commit)
• Before deploying to production
• Never on every test run