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.

The day vi.mock stops being enough — dropping to MSW at the HTTP layer

I said above that MSW becomes the better option as a project matures, so here is what that transition actually feels like.

vi.mock imitates the SDK's class structure. Which means your tests are coupled to the SDK's internal shape. When I upgraded to a version where the return value moved from result.response.text() to result.text, a dozen tests went red at once — while the application behaviour hadn't changed at all. What broke was the mock, not the code.

MSW sits one layer lower and replaces the HTTP request and response. Whatever the SDK does internally, the POST it eventually sends to generativelanguage.googleapis.com keeps the same shape. Pin that, and an SDK upgrade stops dragging your suite down with it.

// tests/msw/handlers.ts
import { http, HttpResponse } from "msw";
 
const ENDPOINT = "https://generativelanguage.googleapis.com/v1beta/models/:model";
 
export const handlers = [
  http.post(`${ENDPOINT}\\:generateContent`, async ({ request }) => {
    const body = (await request.json()) as { contents?: unknown[] };
    if (!Array.isArray(body.contents) || body.contents.length === 0) {
      return HttpResponse.json(
        { error: { code: 400, status: "INVALID_ARGUMENT" } },
        { status: 400 },
      );
    }
    return HttpResponse.json({
      candidates: [
        {
          content: { role: "model", parts: [{ text: "1. Fast delivery" }] },
          finishReason: "STOP",
        },
      ],
      usageMetadata: { promptTokenCount: 32, candidatesTokenCount: 12 },
    });
  }),
];

Streaming adds ?alt=sse and returns a sequence of data: lines. Build the ReadableStream yourself and you get to test your SSE-folding logic too.

http.post(`${ENDPOINT}\\:streamGenerateContent`, () => {
  const chunks = ["Hel", "lo, ", "world"];
  const stream = new ReadableStream({
    start(controller) {
      const enc = new TextEncoder();
      for (const c of chunks) {
        const payload = { candidates: [{ content: { parts: [{ text: c }] } }] };
        controller.enqueue(enc.encode(`data: ${JSON.stringify(payload)}\n\n`));
      }
      controller.close();
    },
  });
  return new HttpResponse(stream, {
    headers: { "Content-Type": "text/event-stream" },
  });
});

One trap worth flagging. Gemini endpoints use a colon separator — models/gemini-2.5-flash:generateContent. In MSW path patterns a colon introduces a path variable, so unless you escape it as \\:generateContent, the handler never matches. Worse, an unmatched request passes through by default, meaning your tests quietly hit the live API. I nearly ran up a bill this way. Once your handlers are written, always start the server with onUnhandledRequest: "error" so a pass-through registers as a failure rather than a green run.

The two approaches protect against different things.

You don't need to move everything to MSW. What I settle on is a two-layer split: branch-heavy tests like retry predicates and tool-loop control stay on vi.mock, and only the three or four tests that guard "what goes into the request body" and "how SSE gets folded" drop down to MSW. Move everything down and the suite slows; leave everything up and you've staked your test suite's lifespan on the SDK's minor versions.

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.

The one thing mocks can't see — a single nightly live call

A mock is a well-constructed lie. Because it's internally consistent, it never tells you the outside world moved. What slipped past me was a model ID going deprecated: the suite stayed green while production quietly failed.

So I keep exactly one test that hits the real API. It doesn't run in PR CI — only on a nightly schedule.

// tests/live.smoke.test.ts
import { describe, it, expect } from "vitest";
import { callGeminiRaw } from "../src/raw";
 
const live = process.env.RUN_LIVE_SMOKE === "1";
 
describe.skipIf(!live)("live smoke", () => {
  it("the configured model still answers", async () => {
    const res = await callGeminiRaw("ping");
    expect(res.candidates?.[0]?.finishReason).toBe("STOP");
    expect(res.candidates?.[0]?.content?.parts?.[0]?.text ?? "").not.toBe("");
    expect(res.usageMetadata?.promptTokenCount).toBeGreaterThan(0);
  }, 30_000);
});

Keeping the assertions deliberately thin is the whole point. It never looks at the content. It checks three things: the model ID is alive, auth still works, and the response skeleton hasn't shifted. Add an expectation about wording and you've built a test that fails most nights on model variance — and one nobody looks at anymore.

If setupFiles overwrites your API key with a dummy, this one test needs the real one. Add a condition to the overwrite:

// tests/setup.ts
if (process.env.RUN_LIVE_SMOKE !== "1") {
  process.env.GOOGLE_GENERATIVE_AI_API_KEY = "test-key-not-real";
}

Cost is one call a day, roughly thirty a month. That fits inside the free tier comfortably, and a few tokens to gemini-2.5-flash never becomes a line item you notice.

Decide up front how you'll treat a failure, too. A red here means "the ground moved," not "my code broke." I route it to a notification I check the next morning rather than something that blocks merges — it makes little sense for an unrelated change to be stuck because the platform shifted overnight. It's the same line I drew earlier when I kept prompt snapshots at warning level.

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.

When you have room for the next step, it's one of two things. If an SDK migration is anywhere on the horizon, drop a handful of tests down to MSW. If your code runs on a nightly batch or any schedule, add the one live smoke test. Either takes an afternoon, and both start paying off the following morning.