How to unit test LlamaIndex query engines

Unit testing LlamaIndex query engines is straightforward once you stop treating the LLM as a black box and start injecting mocks for every network-bound dependency. The pattern below shows how to build a fully deterministic query engine that runs in milliseconds and fails only when your logic is wrong, not when a provider is rate-limited.

Step 1: Pin dependencies and scaffold a pytest suite

Start with a clean virtual environment and pin the core package. LlamaIndex ships frequent breaking changes in minor versions, so a unit test that passes today should still pass next quarter.

python -m venv .venv && source .venv/bin/activate
pip install "llama-index-core==0.10.43" pytest==8.2.0
mkdir tests && touch tests/test_query_engine.py

Keep your test file focused. A good unit test for a query engine imports only the components under test and their mocks—no llama_index.llms.openai imports, no environment variables for API keys.

Step 2: Swap the LLM for a deterministic MockLLM

The fastest way to remove network calls is to replace the LLM with MockLLM. It returns a fixed string (or a programmed sequence) regardless of the prompt, which makes assertions trivial.

from llama_index.core.llms.mock import MockLLM

def test_mock_llm_basic():
    llm = MockLLM(response="The answer is 42")
    out = llm.complete("Any prompt here")
    assert out.text == "The answer is 42"

For query engines that use chat models, MockLLM also implements chat and returns a ChatResponse with the same response string. If you need to simulate streaming, pass streaming=True and iterate the generator—but most unit tests don’t need that path.

Step 3: Stub embeddings with a fixed-vector MockEmbedding

Vector stores need embeddings. MockEmbedding returns a constant vector (or a hash-based one) so similarity search becomes deterministic.

from llama_index.core.embeddings.mock import MockEmbedding

embed = MockEmbedding(embed_dim=8)
# All texts map to the same vector; retrievers return nodes in insertion order.

If your logic depends on semantic ordering, write a tiny custom embedding that encodes a keyword into the vector instead:

from llama_index.core.embeddings.base import BaseEmbedding

class KeywordEmbedding(BaseEmbedding):
    def _get_text_embedding(self, text: str) -> list[float]:
        # Place "urgent" docs at index 0, others at index 1
        return [1.0, 0.0] if "urgent" in text.lower() else [0.0, 1.0]

This level of control is the heart of unit testing LlamaIndex query engines: you decide what “relevant” means.

Step 4: Build an in-memory index from fake documents

With mocks in place, construct a VectorStoreIndex directly from Document objects. No disk, no Chroma, no Pinecone.

from llama_index.core import Document, VectorStoreIndex

docs = [
    Document(text="The ultimate answer is 42.", id_="doc1"),
    Document(text="Python is a snake.", id_="doc2"),
]

index = VectorStoreIndex.from_documents(
    docs,
    embed_model=embed,
    llm=llm,
)
query_engine = index.as_query_engine()

Because MockEmbedding returns identical vectors, the default retriever returns both nodes with equal score. That’s fine for testing the synthesis step. If you used KeywordEmbedding, only the matching doc surfaces.

Step 5: Assert on response, metadata, and source nodes

A query engine returns a Response object containing the synthesized text and the source_nodes that fed it. Test all three.

def test_query_engine_returns_mock():
    resp = query_engine.query("What is the answer?")
    assert resp.response == "The answer is 42"
    assert len(resp.source_nodes) == 2
    assert resp.source_nodes[0].node.id_ == "doc1"

When you practice unit testing LlamaIndex query engines, never assert only on the string. The source_nodes prove the retriever wired up correctly. A wrong retriever that returns empty nodes but a mocked LLM would still pass a string-only test and ship broken.

Step 6: Test custom retrievers and postprocessors in isolation

Real query engines often use a custom BaseRetriever or a NodePostprocessor. Mock the retriever entirely to test the engine’s handling of your postprocessor.

from llama_index.core.retrievers import BaseRetriever
from llama_index.core.schema import NodeWithScore, TextNode, QueryBundle
from llama_index.core.query_engine import RetrieverQueryEngine

class FixedRetriever(BaseRetriever):
    def _retrieve(self, query_bundle: QueryBundle) -> list[NodeWithScore]:
        node = TextNode(text="fixed node", id_="n1")
        return [NodeWithScore(node=node, score=0.9)]

def test_postprocessor_runs():
    retriever = FixedRetriever()
    engine = RetrieverQueryEngine.from_args(
        retriever, llm=llm, node_postprocessors=[]
    )
    resp = engine.query("anything")
    assert resp.source_nodes[0].node.id_ == "n1"

Add a postprocessor that filters low-score nodes and assert it drops them. This isolates your logic from LlamaIndex’s internals.

Step 7: Run the suite and verify success

Execute pytest with verbose output to confirm each test collects and passes.

pytest tests/test_query_engine.py -v

A green suite shows lines like:

tests/test_query_engine.py::test_mock_llm_basic PASSED
tests/test_query_engine.py::test_query_engine_returns_mock PASSED
tests/test_query_engine.py::test_postprocessor_runs PASSED

Total runtime should be under one second. If you see import errors for llama_index.core.llms.mock, check your version—older releases placed MockLLM under llama_index.llms.mock.

Moving to integration checks

The mocks above prove your wiring. For a live smoke test, point the same QueryEngine at a real model via an OpenAI-compatible client. For example, n4n.ai exposes one OpenAI-compatible endpoint covering 240+ models with automatic fallback when a provider is degraded; swapping MockLLM for OpenAI(apis_base=..., api_key=...) requires no changes to index or query engine construction. That step belongs in CI as an optional marker, not in unit tests.

The discipline of unit testing LlamaIndex query engines pays off when you refactor prompts or add rerankers: the tests catch broken composition before users do.