[KongchangAI]
· 3 min read· 1,751 words

Testing LLM Integration Code: A Practical Guide to Mocking the OpenAI Client with pytest

Testing LLM Integration Code: A Practical Guide to Mocking the OpenAI Client with pytest

Use dependency injection and a lightweight fake client to test LLM integration code without hitting the real API.

This article presents an engineering approach to testing code that calls large language models. The core insight is that the logic worth testing — message construction, JSON parsing, error handling, retry behavior — is deterministic and lives around the model, not inside it. By injecting the OpenAI client as a parameter and building a lightweight fake using `SimpleNamespace`, you can run most tests in milliseconds with no network calls. The fake client records every request for assertion and accepts exceptions as scripted replies. A small set of skip-by-default live smoke tests handle the rest, checking only response shape to catch configuration errors.

How do you test code that calls a large language model? It's a challenge many AI feature developers run into. Aman Kumar, who built an OpenAI-compatible API gateway and has spent years reading others' integration code, noticed a recurring anti-pattern: test suites that call the real model on every run. This approach is slow, flaky, burns through your request quota, and when a test fails, you can't tell whether your code is broken or the model just happened to phrase its answer differently today.

This article summarizes his alternative approach. It works with the official openai Python package pointed at any OpenAI-compatible base URL, regardless of which provider you're using.

Core Principle: Tests Should Check Your Code, Not the Model

Most things that go wrong in LLM features aren't in the model itself — they're in the code surrounding it: building the messages list, selecting the right model ID, parsing JSON from the response, handling empty choices lists, retrying on 429s, counting tokens. All of this logic is deterministic and fully testable, as long as you prevent network calls from happening.

Testing LLM integration code: how to mock the OpenAI client in pytest

With that in mind, the author splits tests into two groups:

  1. Unit tests (the vast majority). The client is replaced with a fake, runs in milliseconds, never touches the network.
  2. Live smoke tests (a small handful). These hit the real endpoint, are skipped by default, and only check the shape of the response — not the specific wording.

The value of this split: 99% of your deterministic logic goes to fast, free, reliable unit tests, while a small number of real verifications are left to smoke tests to catch configuration errors that fakes can't detect.

Step 1: Leave a "Seam" for the Client

The hardest code to test is code that creates a client at module import time:


The concept of a "seam" comes from Michael Feathers' definition in *Working Effectively with Legacy Code*: a place in the program where you can alter behavior without editing the code itself. In object-oriented code, the most common form of seam is dependency injection — instead of hardcoding an external dependency (like a network client or database connection) inside a function, you pass it in as a parameter. The immediate benefit: callers can pass in a fake during testing and a real object in production, and the business logic never knows the difference. The anti-pattern is instantiating the client at the module level (`client = OpenAI()`), which turns it into global state. Replacing it in tests then requires `unittest.mock.patch` to monkey-patch the module attribute — fragile and tightly coupled to module paths. Dependency injection is the foundational prerequisite that makes this entire testing approach work.
# Hard to test
from openai import OpenAI
client = OpenAI()

def summarize(text):
    r = client.chat.completions.create(model="gpt-4o-mini", messages=[...])
    return r.choices[0].message.content

The author's approach is to pass the client in as a parameter, with a small factory function for production use:

# llm.py
import os
from openai import OpenAI

def make_client():
    return OpenAI(
        base_url=os.environ.get("LLM_BASE_URL", "https://api.openai.com/v1"),
        api_key=os.environ["LLM_API_KEY"],
        timeout=60,
        max_retries=2,
    )

def summarize(client, text, model):
    r = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": "Summarize in 2 sentences."},
            {"role": "user", "content": text},
        ],
        temperature=0.2,
    )
    if not r.choices:
        raise ValueError("empty response")
    return (r.choices[0].message.content or "").strip()

Two design points worth noting: both the model ID and base URL come from configuration, meaning you can switch providers without touching the business function. And the client is injected as a parameter, which naturally leaves a "seam" for swapping it out in tests.

Step 2: A Minimal Fake Client

You don't need any mocking library to do this. The access path on the OpenAI client is client.chat.completions.create(...), so the fake just needs to provide that path and return an object with .choices[0].message.content.

# tests/fakes.py
from types import SimpleNamespace as NS

def completion(content, finish_reason="stop", prompt_tokens=10, completion_tokens=5):
    msg = NS(content=content, role="assistant", tool_calls=None)
    return NS(
        choices=[NS(index=0, message=msg, finish_reason=finish_reason)],
        usage=NS(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens,
                 total_tokens=prompt_tokens + completion_tokens),
    )

class FakeCompletions:
    def __init__(self, replies):
        self.replies = list(replies)
        self.calls = []
    def create(self, **kwargs):
        self.calls.append(kwargs)
        reply = self.replies.pop(0)
        if isinstance(reply, Exception):
            raise reply
        return reply

class FakeClient:
    def __init__(self, *replies):
        self.chat = NS(completions=FakeCompletions(replies))

Two things make this fake genuinely useful beyond just "returning a string": the calls list records every request, letting you assert what your code actually sent; and replies can be exception objects, letting you script out a sequence of failures. These two features are what elevate the fake from a toy into a tool that catches real bugs.

Step 3: Tests That Actually Catch Bugs

# tests/test_llm.py
import pytest
from llm import summarize
from tests.fakes import FakeClient, completion

def test_sends_expected_request():
    fake = FakeClient(completion("Short summary."))
    out = summarize(fake, "long text", model="my-model")
    assert out == "Short summary."
    call = fake.chat.completions.calls[0]
    assert call["model"] == "my-model"
    assert call["messages"][0]["role"] == "system"
    assert call["messages"][1]["content"] == "long text"

def test_none_content_becomes_empty_string():
    fake = FakeClient(completion(None))
    assert summarize(fake, "x", model="m") == ""

def test_empty_choices_raises():
    fake = FakeClient(completion("x"))
    fake.chat.completions.replies[0].choices = []
    with pytest.raises(ValueError):
        summarize(fake, "x", model="m")

The None content case is a real gotcha: when the model returns only a tool call, or when the response is filtered, message.content can be None. The author admits to seeing production code crash more than once from calling .strip() on None. This is exactly where unit tests earn their keep — they force you to explicitly handle edge cases like this.

Testing Retry Logic and Error Handling

If you wrap calls with your own retry logic, you should script out failure sequences. Real exception classes require a response object, so the author typically raises lightweight stand-in exceptions in unit tests and tests the error-mapping logic separately:

class FakeRateLimit(Exception):
    pass

def call_with_retry(client, attempts=3, **kw):
    for i in range(attempts):
        try:
            return client.chat.completions.create(**kw)
        except FakeRateLimit:
            if i == attempts - 1:
                raise

def test_retries_then_succeeds():
    fake = FakeClient(FakeRateLimit(), FakeRateLimit(), completion("ok"))
    r = call_with_retry(fake, model="m", messages=[])
    assert r.choices[0].message.content == "ok"
    assert len(fake.chat.completions.calls) == 3

In real code, you'd catch openai.RateLimitError and openai.APITimeoutError, and inject the exception type or sleep function so tests don't actually wait out the backoff period.

In the OpenAI Python SDK, both openai.RateLimitError and openai.APITimeoutError inherit from openai.APIStatusError, whose constructor requires a real HTTP response object — making them cumbersome to instantiate directly in unit tests. Two common workarounds: define lightweight stand-in exception classes (like FakeRateLimit) and make the retry logic catch an injectable exception type rather than hardcoding openai.RateLimitError; or use unittest.mock.MagicMock to construct a fake response object and pass it to the real exception. The first approach is cleaner; the second is closer to real behavior. Also, if you use exponential backoff, make time.sleep injectable or patch it with unittest.mock.patch — otherwise unit tests will slow down due to real waits, which defeats the goal of millisecond-level execution.

Testing JSON Parsing

If your feature requires the model to return JSON, the parser is usually where bugs concentrate most. Feed it the kinds of "messy" output models actually produce:

import json, re

def parse_json_reply(text):
    text = text.strip()
    m = re.search(r"```(?:json)?\s*(.*?)```", text, re.S)
    if m:
        text = m.group(1)
    return json.loads(text)

The author's most valuable advice here: collect real error responses from your logs and add them as test cases one by one. That list will become the most precious test file in the entire project — because it captures the real distribution of model behavior, not the idealized output you imagined.

Streaming Responses and Live Smoke Tests

For streaming responses, the fake needs to return an iterable of chunks shaped like chunk.choices[0].delta.content. The key thing to test is that your code correctly concatenates these fragments and handles chunks where delta.content is None (which is common for the first and last chunks).

Since fakes can gradually diverge from reality, the author keeps a small number of real calls that are skipped by default and run on demand:


Streaming responses work like this: when `stream=True`, the OpenAI API returns a Server-Sent Events stream, which the SDK wraps into an iterable that yields `ChatCompletionChunk` objects on each iteration. Unlike non-streaming calls, the full text must be assembled by concatenating `choices[0].delta.content` from each chunk. The first chunk's `delta` typically only contains a `role` field with `content` as `None`; the final chunk has `finish_reason` set to `"stop"` and `content` also as `None`. Streaming code must therefore defensively handle `None` (e.g., `content = chunk.choices[0].delta.content or ""`). Simulating streaming in a fake simply requires `create()` to return a plain Python list or generator whose elements are `SimpleNamespace`-constructed chunk objects — the business logic's `for chunk in stream` loop doesn't need to know the difference.
# tests/test_live.py
import os, pytest
from llm import make_client, summarize

live = pytest.mark.skipif(not os.environ.get("RUN_LIVE"), reason="set RUN_LIVE=1")

@live
def test_live_summary_shape():
    client = make_client()
    out = summarize(client, "The cat sat on the mat. It was warm.",
                    model=os.environ["LLM_MODEL"])
    assert isinstance(out, str) and 0 < len(out) < 500

Notice the assertion: it only checks the type and approximate length — never the exact text. The author only runs these tests before a release or after switching a provider, base URL, or model ID — not on every commit. When pointed at a gateway, this is also how you confirm that a model ID string is correct: a wrong ID typically fails here with a 404 or "model not found" error, rather than surfacing in production.

Anti-Patterns to Avoid

The author explicitly lists several patterns he skips:

  • Snapshot-testing model output. It breaks whenever the model changes, and gives you zero feedback about code quality.
  • Recording real responses with cassette tools for every test. It works, but recordings go stale and mask changes in response shape. The author only uses recording for a handful of complex tool-calling flows.
  • Mocking requests/httpx at the SDK's transport layer. This couples your tests to SDK internals. The right boundary to mock at is the client interface.

Summary

The whole approach boils down to a few sentences: pass the client in as a parameter, fake out chat.completions.create, assert both what you received and what you sent, script out failure sequences, and keep a small set of on-demand live tests that only check shape. Your test suite becomes fast, deterministic, and free to run — while the handful of live tests catch configuration errors that fakes can't.

For any team integrating LLMs in production, this layered approach of "mock at the client boundary + shape-only smoke tests" is far more robust than blindly calling real models, and it's an engineering practice well worth adopting.

#mock OpenAI client#pytest LLM testing#test LLM integration code#fake OpenAI client#unit test AI code#OpenAI dependency injection#LLM smoke tests
Share:

Related articles