Mocking the Vector Store to Test Generation Logic in Isolation

Half of a retrieval system is code that turns a list of chunks into a prompt: ordering them, labelling them, cutting them to fit, deciding what to say when there are none. That half is deterministic, and it can be tested completely without an index or a model.

The seam, and what sits on each side

Put the boundary at a function that takes a query and returns chunks. Above it: embedding the query, the index, filters, reranking. Below it: ordering, deduplication, token budgeting, template assembly, citation labelling. This page is about the second list, and mocking the first is how you get at it.

type Chunk = {
  ref: string;          // stable chunk identifier
  doc: string;          // source document
  text: string;
  score: number;
  meta: { title: string; url: string; updated: string };
};

type Retrieve = (query: string, k: number) => Promise<Chunk[]>;

export function buildPrompt(query: string, chunks: Chunk[], budget: number): {
  messages: Msg[];
  used: string[];       // chunk refs that made it in
  dropped: string[];    // chunk refs cut for budget
} { /* ... */ }

Returning used and dropped alongside the messages is what makes this testable without parsing the prompt string back apart. It is also useful in production for exactly the same reason: when an answer misses something that was retrieved, the question is always whether the chunk was dropped, and a function that only returns a string cannot answer it.

Inject the retriever rather than importing it. A test that has to reach past a module boundary to control what comes back ends up mocking the index client, which pulls the index client’s own quirks into every test below the seam. A function-typed parameter with a real implementation as its default keeps production code unchanged and makes the test a one-line substitution.

The prompt is the output under test

Assert on the messages, not on any answer. There is no model in this test at all.

import { describe, it, expect } from "vitest";
import { chunk } from "./fixtures/chunks";   // small builder with defaults

it("includes every chunk, in rank order, with its citation label", () => {
  const chunks = [
    chunk({ ref: "a1", doc: "refunds", text: "Refunds take 5 business days." }),
    chunk({ ref: "b2", doc: "returns", text: "Post returns within 30 days." }),
  ];

  const { messages, used, dropped } = buildPrompt("how long?", chunks, 4000);

  expect(used).toEqual(["a1", "b2"]);
  expect(dropped).toEqual([]);

  const context = messages.find((m) => m.role === "system")!.content as string;
  expect(context.indexOf("Refunds take")).toBeLessThan(
    context.indexOf("Post returns"),
  );
  expect(context).toContain("[1]");
  expect(context).toContain("[2]");
  expect(context).not.toContain("undefined");
});

The not.toContain("undefined") line looks like a joke and is the assertion most likely to fire. Templates interpolate metadata fields, real chunks are missing metadata more often than anyone expects, and the string undefined — or None, in Python — appearing in a prompt is a bug the model papers over by ignoring it. Assert its absence once and you catch every future field that goes missing.

Assert on ordering by index position rather than on the exact assembled string. A whole-string assertion breaks every time somebody adjusts a heading, which teaches people to update the expectation without reading it. Positional assertions break only when the order changes, which is the thing you meant.

The context states worth forcing

This is the real reason to mock the store. Each of these happens in production and none can be produced on demand from a real index.

• No chunks at all. The retriever returned an empty list. The prompt must say so explicitly and instruct the model not to answer from memory — a template that simply interpolates an empty context block produces a fluent, unsourced, confident answer, which is the worst output a retrieval system can produce. Assert that the no-context branch was taken, not merely that the string is short.
• One weak chunk. A single result well below your score floor. Whether you pass it through or treat it as nothing is a product decision; whichever you chose, pin it.
• Contradictory chunks. Two documents saying 5 days and 10 days. You cannot assert what the model does, but you can assert that both survived into the prompt with distinguishable citations and dates — a pipeline that silently deduplicates them by similarity leaves the model no way to notice the conflict.
• Near-duplicate chunks. The same paragraph indexed twice from two documents. Assert your deduplication kept one, and kept the higher-ranked one.
• A chunk containing an instruction. Retrieved text saying “ignore previous instructions”. Assert on how it is delimited and labelled as data, which is the structural half of injection defence. You cannot test that the model resists it; you can test that your template did not hand it over undelimited.

What happens when the context does not fit

Budgeting is the most consequential code below the seam and the least examined. Fixed context makes it exactly testable: build chunks of known length, set a budget that admits three of five, and assert which two were dropped.

it("drops the lowest-ranked chunks first and reports them", () => {
  const chunks = ["a", "b", "c", "d", "e"].map((r, i) =>
    chunk({ ref: r, text: "x".repeat(400), score: 1 - i * 0.1 }),
  );

  const { used, dropped } = buildPrompt("q", chunks, 300);

  expect(used).toEqual(["a", "b", "c"]);
  expect(dropped).toEqual(["d", "e"]);
});

Two properties matter more than the exact numbers. That dropping is by whole chunks, not by truncating the last one mid-sentence — a half chunk is a half fact, and it is worse than no chunk because the model cannot tell it is incomplete. And that dropping is from the bottom of the ranking, which is only true if the budget is applied after ordering; applying it during retrieval instead means the highest-scoring chunk can be the one that does not fit.

Test the pathological case too: one chunk larger than the entire budget. The function must not return an empty prompt and must not loop; deciding between “truncate this one chunk with a marker” and “fail loudly” is a real choice, and an untested implementation usually does neither on purpose.

Budget in the same unit production budgets in. If the ceiling is a token count, the test must count tokens, because characters and tokens diverge badly on code, on non-Latin scripts and on anything with long identifiers — a chunk that is 400 characters can be 90 tokens of prose or 200 of JSON. A character-based test against a token-based budget passes while production truncates, which is the worst arrangement available. If counting tokens in a unit test is too slow, assert on characters with a deliberately conservative ceiling and say in the test name that it is a proxy.

What this test cannot tell you

It says nothing about whether the retrieved chunks were the right ones — you supplied them — and nothing about whether the model uses them well. Those are the retriever’s own tests and an evaluation, respectively.

Its value is that when an answer is wrong, it removes an entire half of the system from the investigation. If the assembly tests pass and the retriever tests pass, the fault is in retrieval quality or in the model, and you have narrowed it in seconds rather than by staring at a prompt. That is worth more than it sounds: the alternative involves reading several thousand tokens of assembled context looking for something that is not there. Keep the fixed chunks in the same file as the tests, small enough to read; the moment they become a large shared corpus you have rebuilt the thing you were trying to mock out.