Testing Gemini API with Vitest — Mock Patterns and Function Calling Verification for TypeScript Projects

Setup — the smallest config that works

Install Vitest and the SDK:

npm install -D vitest @vitest/coverage-v8
npm install @google/generative-ai

The configuration can stay tiny:

// vitest.config.ts
import { defineConfig } from "vitest/config";
 
export default defineConfig({
  test: {
    environment: "node",
    globals: true,
    coverage: { provider: "v8" },
  },
});

To stop tests from accidentally hitting the real Gemini API, I overwrite the API key in a setup file:

// tests/setup.ts
process.env.GOOGLE_GENERATIVE_AI_API_KEY = "test-key-not-real";

Add setupFiles: ["tests/setup.ts"] under test in the config. Even if your .env has a real key, the setup file wins, so you can run the suite in CI without burning quota.

Mocking the SDK module — the first test you should write

The first thing I want to lock down isn't the response text. It's that my wrapper function calls the SDK with the right arguments. The contents of the response are the model's responsibility; the call site is mine.

vi.mock at the module level handles this cleanly. Here's a thin wrapper around Gemini:

// src/summarize.ts
import { GoogleGenerativeAI } from "@google/generative-ai";
 
export async function summarizeReview(text: string) {
  const genAI = new GoogleGenerativeAI(process.env.GOOGLE_GENERATIVE_AI_API_KEY!);
  const model = genAI.getGenerativeModel({ model: "gemini-2.5-flash" });
  const result = await model.generateContent({
    contents: [{ role: "user", parts: [{ text }] }],
    generationConfig: { temperature: 0.2, maxOutputTokens: 200 },
  });
  return result.response.text();
}

And the test:

// tests/summarize.test.ts
import { describe, it, expect, vi, beforeEach } from "vitest";
 
const generateContent = vi.fn();
vi.mock("@google/generative-ai", () => {
  return {
    GoogleGenerativeAI: vi.fn().mockImplementation(() => ({
      getGenerativeModel: () => ({ generateContent }),
    })),
  };
});
 
import { summarizeReview } from "../src/summarize";
 
beforeEach(() => {
  generateContent.mockReset();
});
 
describe("summarizeReview", () => {
  it("passes temperature 0.2 and the user text correctly", async () => {
    generateContent.mockResolvedValue({
      response: { text: () => "1. Fast shipping\n2. Matches description\n3. Good value" },
    });
 
    const out = await summarizeReview("The product was excellent.");
 
    expect(generateContent).toHaveBeenCalledTimes(1);
    const arg = generateContent.mock.calls[0][0];
    expect(arg.generationConfig.temperature).toBe(0.2);
    expect(arg.contents[0].parts[0].text).toContain("excellent");
    expect(out).toMatch(/shipping/);
  });
});

The key discipline: don't pin the response text with strict equality. The model decides the wording, so use toMatch for keywords you'd reasonably expect. Pin the input — the call shape — and snapshot it from there. That's where most of the long-term value is.

Testing Function Calling — fixing the contract, not the model's choice

When you write Function-Calling code, it's tempting to assert that "the model picked tool X." Don't. That depends on the model, and your tests will flake. I split it into two layers instead: was the tool definition assembled correctly, and once the tool returns, does my loop continue as expected.

The definition side, with no model calls at all:

// tests/tools.test.ts
import { describe, it, expect } from "vitest";
import { buildSearchTools } from "../src/tools";
 
describe("buildSearchTools", () => {
  it("includes 'query' as a required parameter on the search tool", () => {
    const tools = buildSearchTools();
    const fn = tools[0].functionDeclarations[0];
    expect(fn.name).toBe("search_products");
    expect(fn.parameters.required).toContain("query");
  });
});

The continuation side uses chained mockResolvedValueOnce calls. Function Calling in Gemini comes back as a "use this tool" response, then a follow-up response after you provide the tool result, so two stages is the natural shape:

generateContent
  .mockResolvedValueOnce({
    response: {
      functionCalls: () => [{ name: "search_products", args: { query: "tshirt" } }],
      text: () => "",
    },
  })
  .mockResolvedValueOnce({
    response: { text: () => "Recommended: Tshirt A, B, C.", functionCalls: () => [] },
  });

This way, even if the model changes its mind in production, the tool assembly and your control flow remain pinned. When something breaks, you can tell prompt drift apart from a code bug — which is the whole reason we're writing these tests.

Streaming responses — mocking AsyncIterable

For code that uses generateContentStream, the mock needs to return an AsyncIterable. I tripped on this the first time, so here's the template I keep around:

async function* fakeStream(chunks: string[]) {
  for (const c of chunks) {
    yield { text: () => c };
  }
}
 
generateContentStream.mockResolvedValue({
  stream: fakeStream(["Hello", ", ", "world"]),
});

Your call site presumably consumes this with for await and concatenates. Assert that the joined output is "Hello, world". To test mid-stream errors, define a generator that throws partway through and inject that instead — your error-handling path gets coverage for free.

Pin the prompt as a contract using snapshots

The cheapest regression check for prompts is expect(prompt).toMatchSnapshot(). Treat the prompt string itself as a versioned asset.

import { describe, it, expect } from "vitest";
import { buildReviewPrompt } from "../src/prompt";
 
describe("buildReviewPrompt", () => {
  it("freezes the shape of the review-summary prompt", () => {
    const p = buildReviewPrompt({ productName: "Tshirt A", review: "loved it" });
    expect(p).toMatchSnapshot();
  });
});

When you intentionally change the prompt, run vitest -u to update the snapshot. When you didn't intend to, the diff shows up in the PR, and the reviewer catches it. It also raises the visibility of prompts as code: they're not just strings to twiddle, they're contracts.

If you want to push beyond shape into output quality, the Pytest-side article Don't break your Gemini prompts: regression testing patterns with Pytest covers an LLM-as-Judge layer over the same idea. The two-tier approach maps directly onto Vitest.

Should it fail CI, or just warn? — the practical split

A small operational note. I treat code-side tests (tool assembly, streaming concatenation, error handling) as hard CI gates. They're deterministic, so red means broken, full stop.

Snapshot tests on prompts I run as warnings only for the first few weeks. Updating the snapshot is on the PR author and reviewer, not the CI gate. That separates "I want to iterate on this prompt" from "I want to merge today," which keeps prompt refactors from feeling expensive. As things stabilize, tighten gradually.

The same gradual approach applies to the mock layer. Module-level mocking via vi.mock (the setup in this article) is the right starting point. As production traffic gets serious, MSW at the HTTP layer becomes more durable across SDK upgrades. Don't pay for that complexity until the project is around to benefit.

A note on test data — keep fixtures small and named

One small habit that's paid off for me: store test inputs as named fixtures, not inline strings. After the second or third test that involves "a Japanese review with a typo" or "a review that mentions a delivery problem," it gets hard to remember which one you're looking at when a snapshot diffs.

// tests/fixtures/reviews.ts
export const reviewWithDeliveryComplaint =
  "Item is fine, but it took two weeks to arrive even though shipping was paid extra.";
 
export const reviewWithMixedSentiment =
  "Material is much better than expected. The color is way off from the photos though.";

Importing these by name into tests makes the intent of each assertion legible at a glance, which is when you're debugging a flaky test at the end of the day. It also makes it easier to grow a small bank of "interesting cases" — edge cases I want every new prompt revision to be tested against — over time.

The same pattern applies on the mock side. If you have several scenarios where the model "calls a tool, then returns text," extract the helper:

function mockToolThenAnswer(toolName: string, args: object, finalText: string) {
  generateContent
    .mockResolvedValueOnce({
      response: { functionCalls: () => [{ name: toolName, args }], text: () => "" },
    })
    .mockResolvedValueOnce({
      response: { text: () => finalText, functionCalls: () => [] },
    });
}

Three or four of these helpers cover most of what I write in day-to-day Gemini tests. They're not glamorous, but they make the test suite a place future-me actually wants to come back to.

Test rate limits and timeouts on purpose

The scariest part in production isn't the happy path — it's what happens when a 429 (rate limit) or 503 (transient overload) comes back. In my own setup, code that doesn't pass here doesn't get to ride on an overnight batch. Because the generateContent mock can reject as well as resolve, you can test retry control deterministically.

Start by rejecting once with an error that mimics what the SDK throws, then resolving on the second call:

it("retries exactly once after a single 429 and succeeds", async () => {
  const err = Object.assign(new Error("Resource exhausted"), { status: 429 });
  generateContent
    .mockRejectedValueOnce(err)
    .mockResolvedValueOnce({ response: { text: () => "ok" } });
 
  const out = await summarizeReviewWithRetry("test input");
 
  expect(generateContent).toHaveBeenCalledTimes(2);
  expect(out).toBe("ok");
});

Pin the retry ceiling too. Infinite retries only make rate limiting worse, so guarantee in a test that you give up after three attempts and rethrow:

it("gives up and throws after three failures", async () => {
  const err = Object.assign(new Error("Service unavailable"), { status: 503 });
  generateContent.mockRejectedValue(err);
 
  await expect(summarizeReviewWithRetry("x")).rejects.toThrow(/unavailable/);
  expect(generateContent).toHaveBeenCalledTimes(3);
});

Timeouts pair well with vi.useFakeTimers() so you don't burn real seconds. If you use exponential backoff, advance the clock with await vi.advanceTimersByTimeAsync(2000) to skip a 2-second sleep instantly. Waiting on real time in backoff tests adds tens of seconds to the whole suite, so I always switch to fake timers here.

Here's the checklist I actually keep for this:

1. One failure then a retry that succeeds (call count 2)
2. Failures up to the ceiling, throwing and calling no further (call count = ceiling)
3. Anything other than 429/503 (e.g. a 400 malformed request) throws immediately, with no retry
4. The backoff wait can be skipped with fake timers

Item 3 is the easy one to miss. If a prompt-assembly mistake is returning a 400 and you dutifully retry it, you delay finding the bug. Pinning the "is this error retryable?" predicate as its own unit test is worth the few lines.

How much faster did it get — measured on my own setup

One number worth leaving here. The review-summarization tests I run in my own indie developer projects currently sit at 41 cases. Running all of them against mocks, never touching real Gemini once, finishes in roughly 0.9 seconds on my machine. Running the same 41 against the live API would cost 1.5–3 seconds each on average (model and prompt length dependent), so call it 60–120 seconds. Once you're running this on every push in CI, that gap stops being ignorable.

Speed isn't the only win. Because nothing hits the real API, CI never spends tokens. At 41 cases times ten pushes a day, that's over ten thousand API calls a month avoided entirely. For a free-tier solo project, "the tests never touch billing or rate limits" translates straight into peace of mind.

On gathering the numbers: Vitest prints per-file timing by default, but vitest run --reporter=verbose is easier to read for the whole picture. When I find a slow test, it's almost always one of two things — a backoff wait I forgot to fake, or an oversized fixture. The former goes away with vi.useFakeTimers() as above; the latter, by trimming the fixture to the minimum, and most cases drop back under a second.

Start by pinning the call shape, nothing more

Thanks for reading this far. Tests can grow without limit, so the version I keep recommending for solo projects is small: write one test that fixes "what arguments does my code pass to generateContent," and stop there for now. With that pinned, when you iterate on prompts, you can always tell whether a regression is in your code or in the prompt itself — and that single split saves more time than any other test you'll add.