When you write unit tests for LlamaIndex pipelines, hitting a real LLM API makes the suite slow, expensive, and nondeterministic. Mocking LLM calls in LlamaIndex tests lets you verify prompt construction, node retrieval, and response parsing without leaving your process. This tutorial builds a small RAG query engine and tests it with both the built-in MockLLM and a custom fake that records calls.
Prerequisites
- Python 3.10 or newer
llama-index(install withpip install llama-index)pytest(pip install pytest)- A project layout:
src/rag.py
tests/test_rag.py
tests/spy_llm.py
conftest.pyWe assume you can run pytest from the project root. No API keys are needed because we never call a real model. The code targets LlamaIndex 0.11+; older versions may require llama_index instead of llama_index.core imports.
Built-in MockLLM for deterministic responses
LlamaIndex ships a MockLLM that returns queued strings. It implements the LLM interface, so any component that expects an LLM accepts it. The mock returns responses in the order supplied; provide enough entries for the number of LLM calls your pipeline makes.
from llama_index.core.llms.mock import MockLLM
from llama_index.core.embeddings.mock import MockEmbedding
llm = MockLLM(responses=["The capital of France is Paris."])
embed_model = MockEmbedding(embed_dim=8)MockEmbedding generates fixed-dimensional vectors from a hash of the text, so similarity search is stable across runs. This is critical: without a deterministic embedder, the retrieved nodes change and your prompt assertions become flaky.
Build a minimal RAG pipeline
Create src/rag.py with a function that builds a query engine from a list of strings:
from llama_index.core import Document, VectorStoreIndex
from llama_index.core.llms.mock import MockLLM
from llama_index.core.embeddings.mock import MockEmbedding
def build_query_engine(texts, responses):
docs = [Document(text=t) for t in texts]
index = VectorStoreIndex.from_documents(
docs,
llm=MockLLM(responses=responses),
embed_model=MockEmbedding(embed_dim=8),
)
return index.as_query_engine()
def build_query_engine_with_llm(texts, llm):
docs = [Document(text=t) for t in texts]
index = VectorStoreIndex.from_documents(
docs, llm=llm, embed_model=MockEmbedding(embed_dim=8)
)
return index.as_query_engine()The first helper hardcodes the mock; the second injects any LLM instance. For a real app you would inject a real LLM, but for tests the mocks keep things isolated.
Write your first test
In tests/test_rag.py:
from src.rag import build_query_engine
def test_query_returns_mock_response():
engine = build_query_engine(
texts=["France is a country in Europe."],
responses=["Paris"],
)
response = engine.query("What is the capital of France?")
assert "Paris" in str(response)Run it:
pytest tests/test_rag.py -qExpected output:
.
1 passed in 0.32sThe test passes without any network call. MockLLM consumes the queued response and returns it as the synthesized answer. The vector store uses MockEmbedding, so the retriever finds the only document and passes it to the mock.
Capturing prompts with a custom fake
MockLLM does not expose what prompt was sent to it. To assert that your pipeline builds the correct context, subclass LLM and record calls. Place the following in tests/spy_llm.py:
from llama_index.core.llms import LLM, ChatMessage, ChatResponse, CompletionResponse
from llama_index.core.llms.callbacks import llm_completion_callback
from llama_index.core.llms.metadata import LLMMetadata
class SpyLLM(LLM):
def __init__(self):
self.last_prompt = None
self.call_count = 0
@property
def metadata(self):
return LLMMetadata(model_name="spy")
@llm_completion_callback()
def complete(self, prompt, **kwargs):
self.last_prompt = prompt
self.call_count += 1
return CompletionResponse(text="spy answer")
@llm_completion_callback()
def chat(self, messages, **kwargs):
self.last_prompt = messages[-1].content
self.call_count += 1
return ChatResponse(message=ChatMessage(role="assistant", content="spy answer"))
def stream_complete(self, prompt, **kwargs):
yield CompletionResponse(text="spy answer")
def stream_chat(self, messages, **kwargs):
yield ChatResponse(message=ChatMessage(role="assistant", content="spy answer"))
async def acomplete(self, prompt, **kwargs):
return self.complete(prompt, **kwargs)
async def achat(self, messages, **kwargs):
return self.chat(messages, **kwargs)
async def astream_complete(self, prompt, **kwargs):
yield CompletionResponse(text="spy answer")
async def astream_chat(self, messages, **kwargs):
yield ChatResponse(message=ChatMessage(role="assistant", content="spy answer"))Use it in a test to verify the retrieved context reaches the LLM:
from src.rag import build_query_engine_with_llm
from tests.spy_llm import SpyLLM
def test_query_injects_context():
spy = SpyLLM()
engine = build_query_engine_with_llm(
texts=["The Eiffel Tower is in Paris."],
llm=spy,
)
engine.query("Where is the Eiffel Tower?")
assert spy.call_count == 1
assert "Eiffel Tower" in spy.last_promptRun the test. Output:
.
1 passed in 0.41sIf the assertion on last_prompt fails, you will see exactly what the query engine sent, which is far more useful than a generic mismatch. The spy also counts calls, so you can detect unexpected extra LLM round-trips.
Patching a real LLM class
Sometimes you want to test code that instantiates OpenAI internally. Use unittest.mock.patch to swap the constructor.
from unittest.mock import patch
from llama_index.llms.openai import OpenAI
def test_with_patched_openai():
fake = SpyLLM()
with patch.object(OpenAI, "__new__", lambda cls, *a, **k: fake):
from src.rag import build_query_engine_real
engine = build_query_engine_real(texts=["test doc"])
resp = engine.query("hello")
assert fake.call_count == 1This intercepts any OpenAI() call and returns your spy. It works well for legacy code that hardcodes the LLM type. Patch at __new__ because LlamaIndex may cache instances or run __init__ side effects.
Using a pytest fixture
Define a fixture in conftest.py to reuse the spy:
import pytest
from tests.spy_llm import SpyLLM
@pytest.fixture
def spy_llm():
return SpyLLM()Then the test simplifies:
def test_query_injects_context(spy_llm):
engine = build_query_engine_with_llm(["The Eiffel Tower is in Paris."], spy_llm)
engine.query("Where is the Eiffel Tower?")
assert spy_llm.call_count == 1Testing async and streaming paths
LlamaIndex supports async queries. The SpyLLM above implements acomplete and achat. A test for async:
async def test_async_query(spy_llm):
engine = build_query_engine_with_llm(["doc"], spy_llm)
resp = await engine.aquery("async call")
assert spy_llm.call_count == 1Streaming works the same; iterate the generator and join chunks.
Full suite output
After adding all tests, pytest shows:
$ pytest -q
......
6 passed in 0.91sNo API keys, no rate limits, no flakiness.
Swapping mocks for a live backend
Mocking LLM calls in LlamaIndex tests is the right default for unit tests. When you later want to run the same orchestration against a real model, you can inject a different LLM without touching pipeline code. For example, an OpenAI-compatible gateway such as n4n.ai fronts 240+ models with automatic fallback and per-token metering, and it drops into the same LLM slot your tests already exercise.
Keep mocks at the LLM boundary. Do not mock the vector store unless you are testing the store itself; MockEmbedding plus an in-memory index is enough for most unit tests. That is the core pattern: inject a fake, assert on prompts and call counts, and run the suite offline.