TechByteByByte

Unit Testing and Mock Testing

Learn how to write automated tests for AI code using pytest, fixtures, and parametrized tests, and how to mock LLM API calls with unittest.mock so tests run fast, free, and reliably.

#Python#Testing#Pytest#Mocking#AI#Python for AI

The problem: AI applications combine exact Python logic with slow, costly, and variable model calls. Calling a live LLM whenever you check a prompt-building function makes tests fragile and mixes two different kinds of quality.

What you will learn: Unit tests check deterministic logic; evaluations measure open-ended model quality across examples. You will use pytest, fixtures, parametrization, patching, and mocks to replace an API boundary in fast tests. You will also see why a mock proves your code’s reaction to a fake response—not that the real provider or model quality works.


1. Why Testing Matters

A test compares what actually happened with an expectation chosen before the test ran:

arrange known input and dependencies

act by running one behaviour

assert an observable result
        ├── matches → pass
        └── differs → fail with evidence

Tests do not prove that a program has no bugs. They provide repeatable evidence for the cases you selected. Good test cases therefore include normal input, boundaries, invalid input, and known failure paths—not just one happy example.

For AI systems, deterministic Python behaviour and open-ended model quality are different concerns. Unit tests check the former; evaluations over many examples measure the latter.

Code That Checks Code

A test is code that checks your code — it runs a function with known inputs and verifies the output is what you expect, automatically, every time you make a change.

Why Repeatable Checks Matter

Without tests, you find out something broke only when a user hits the bug — or worse, never find out at all. Tests catch breakage the moment it happens, before it reaches anyone.

Picture a Safety Net

Tests are a safety net under a tightrope walker. You can still make bold changes (refactor a function, upgrade a library) because if you slip, the net catches you immediately — a loud, clear failure — instead of the fall happening silently, discovered only much later.

The Pre-Flight Checklist Analogy

Think of a pilot’s pre-flight checklist. It’s run automatically, every single time, regardless of how experienced the pilot is or how many times they’ve flown this exact plane — because manual “I’m sure it’s fine” checking is exactly where mistakes slip through.

🤖 How Is This Used in AI? AI pipelines have a lot of moving, easily-broken pieces — prompt templates, JSON parsing, retry logic, validation. A small change to a prompt-building function can silently break your entire pipeline’s output shape. Tests catch that in seconds, not after a user complains that answers look wrong.


2. Unit Testing Basics — assert

What Is It?

The simplest possible test: assert checks that a condition is True, and raises an error immediately if it isn’t.

def clean_text(text):
    return text.strip().lower()

# A very simple, manual "test"
result = clean_text("  Hello WORLD  ")
assert result == "hello world"
print("Test passed!")

Expected Output:

Test passed!

If the function were broken:

def clean_text(text):
    return text.strip()   # bug: forgot to lowercase

result = clean_text("  Hello WORLD  ")
assert result == "hello world"
AssertionError

🧠 Intuition

assert is a tripwire — silent as long as everything is correct, loudly failing the instant reality doesn’t match your expectation.


3. pytest — Writing Real Tests

What Is It?

pytest is the standard Python testing framework — it finds and runs functions named test_*, reports which passed/failed, and gives clear failure messages.

pip install pytest

text_utils.py

def clean_text(text):
    return text.strip().lower()

def word_count(text):
    return len(text.split())

test_text_utils.py

from text_utils import clean_text, word_count

def test_clean_text_strips_and_lowercases():
    assert clean_text("  Hello WORLD  ") == "hello world"

def test_clean_text_handles_already_clean_input():
    assert clean_text("already clean") == "already clean"

def test_word_count():
    assert word_count("python is great for ai") == 5
pytest test_text_utils.py -v

Expected Output (approximate):

test_text_utils.py::test_clean_text_strips_and_lowercases PASSED
test_text_utils.py::test_clean_text_handles_already_clean_input PASSED
test_text_utils.py::test_word_count PASSED

3 passed in 0.01s

How It Works

  • pytest automatically discovers any file named test_*.py and any function inside it named test_*.
  • Each test function runs independently; assert inside it is how you declare “this must be true.”
  • If an assert fails, pytest reports exactly which test failed, the expected vs. actual values, and the line number — far more useful than a bare AssertionError.

🤖 How Is This Used in AI? Every function from earlier modules — clean_text, build_prompt, cosine_similarity, parse_response — is exactly the kind of pure, predictable logic that deserves a fast, reliable unit test, completely independent of whether any AI API is even reachable.


4. Testing Functions That Raise Exceptions

Recall Module 6’s exceptions — you can test that a function correctly raises an error under bad input:

import pytest

def validate_temperature(temperature):
    if not (0.0 <= temperature <= 2.0):
        raise ValueError("temperature must be between 0.0 and 2.0")
    return temperature

def test_validate_temperature_accepts_valid_value():
    assert validate_temperature(0.7) == 0.7

def test_validate_temperature_rejects_out_of_range():
    with pytest.raises(ValueError):
        validate_temperature(3.5)

🧠 Intuition: pytest.raises(ValueError) says “I expect this block to raise a ValueError — the test should fail if it doesn’t.” This is how you verify your validation logic actually rejects bad input, not just that it accepts good input.

🤖 How Is This Used in AI? Confirming that invalid config (a bad temperature, a negative max_tokens, an empty prompt) is rejected before it ever reaches an expensive API call — exactly the validation logic from Module 6, now with a test proving it works.


5. Fixtures

What Is It?

A fixture is reusable setup code that multiple tests can share — avoiding copy-pasted setup at the top of every test function.

import pytest

@pytest.fixture
def sample_documents():
    return [
        {"text": "Python is great for AI.", "score": 0.9},
        {"text": "Bananas are yellow.", "score": 0.1},
    ]

def test_filters_relevant_documents(sample_documents):
    relevant = [d for d in sample_documents if d["score"] >= 0.5]
    assert len(relevant) == 1
    assert relevant[0]["text"] == "Python is great for AI."

def test_document_count(sample_documents):
    assert len(sample_documents) == 2

🧠 Intuition: A fixture is a shared prop department for a play — instead of every scene (test) building its own props from scratch, they all request the same ready-made prop (sample_documents) from a common source.

🤖 How Is This Used in AI? A fixture is perfect for a reusable “sample chat history,” “sample retrieved documents,” or “sample API response JSON” that many different tests in your AI pipeline need.


6. Parametrized Tests

What Is It?

Running the same test logic against many different input/output pairs, without writing a separate function for each.

import pytest

def score_relevance(similarity):
    if similarity >= 0.8:
        return "high"
    elif similarity >= 0.5:
        return "medium"
    return "low"

@pytest.mark.parametrize("score, expected", [
    (0.9, "high"),
    (0.6, "medium"),
    (0.2, "low"),
    (0.8, "high"),   # boundary case
    (0.5, "medium"),  # boundary case
])
def test_score_relevance(score, expected):
    assert score_relevance(score) == expected

Expected Output:

5 passed in 0.01s

🧠 Intuition

Parametrization is a loop over test cases, but each case is reported individually — if case #4 fails, pytest tells you exactly which input/expected pair broke, instead of one vague failing test covering everything.

🤖 How Is This Used in AI? Testing a relevance-labeling function, a prompt-length truncation function, or a retry-decision function against many realistic boundary values (score thresholds, token limits) all at once — this is one of the most valuable testing patterns for AI logic full of thresholds and edge cases.


7. Mocking — What and Why

What Is It?

A mock is a fake, controllable stand-in for something real — most often, an external API call — used in tests so you don’t actually make a network request.

Why Does It Exist?

Calling a real LLM API in every test run would be:

  • Slow — network calls take real time, tests should run in milliseconds
  • Costly — every test run would spend real API credits
  • Unreliable — network issues or provider outages would fail your tests for reasons that have nothing to do with your code
  • Non-deterministic — an LLM might phrase its answer differently each time, making assert comparisons unreliable

🧠 Intuition

A mock is a flight simulator, not a real plane. You practice and test your reactions to specific scenarios (engine failure, bad weather) without the cost or risk of doing it in a real aircraft — and you can force exact scenarios on demand (“simulate a rate limit error right now”), which is often impossible to reliably trigger with the real thing.

Real-World Analogy

Testing a smoke detector with test-smoke instead of setting an actual fire — you verify the detector’s logic (“does it alarm when smoke is present?”) without needing the real, costly, risky trigger event.

Here is how a mock intercepts outgoing API requests, sealing off your code from the external network during a test:

sequenceDiagram
    participant Test as Test Suite
    participant App as App Function (call_llm)
    participant Mock as Mock Client (fake SDK)
    participant API as Real API Server (Internet)

    Test->>App: call_llm(mock_client)
    App->>Mock: client.messages.create(...)
    Note over Mock: Intercepted!<br/>Return configured stub value
    Mock-->>App: {"text": "Paris is the capital."}
    Note right of API: Real network is blocked!<br/>Cost = $0, Latency = 1ms
    App-->>Test: Return "Paris is the capital."
    Test->>Test: Assert result == "Paris is the capital."

[!NOTE] Mocks vs. Stubs vs. Fakes

  • Mock: A dynamic placeholder configured to return specific values and verify if/how it was called (e.g., mock_client.assert_called_once()).
  • Stub: A simple hardcoded stand-in that returns fixed responses but does not track call counts or arguments.
  • Fake: A working but lightweight implementation suitable for local tests (like using a local in-memory SQLite database instead of a remote production Postgres server).

8. unittest.mock Basics

from unittest.mock import Mock

# A Mock object can pretend to be anything and return whatever you tell it to
fake_client = Mock()
fake_client.messages.create.return_value = {
    "content": [{"text": "This is a fake response."}]
}

response = fake_client.messages.create(model="claude-sonnet-4-6", messages=[])
print(response["content"][0]["text"])

# You can also verify HOW the mock was called
fake_client.messages.create.assert_called_once()
print("Confirmed: create() was called exactly once")

Expected Output:

This is a fake response.
Confirmed: create() was called exactly once

🧠 Intuition

A Mock is an empty stunt double — it will pretend to be any object, respond however you configure it to, and remember every interaction so you can later check “was this actually called, and with what arguments?”


9. Mocking API Calls — the essential AI testing pattern

from unittest.mock import patch

def call_llm(client, prompt):
    """Real function that would call an actual API in production."""
    response = client.messages.create(
        model="claude-sonnet-4-6",
        max_tokens=300,
        messages=[{"role": "user", "content": prompt}],
    )
    return response.content[0].text

def test_call_llm_returns_model_text():
    # Build a fake client shaped like the real SDK's response
    fake_response = Mock()
    fake_response.content = [Mock(text="Paris is the capital of France.")]

    fake_client = Mock()
    fake_client.messages.create.return_value = fake_response

    result = call_llm(fake_client, "What is the capital of France?")

    assert result == "Paris is the capital of France."
    fake_client.messages.create.assert_called_once_with(
        model="claude-sonnet-4-6",
        max_tokens=300,
        messages=[{"role": "user", "content": "What is the capital of France?"}],
    )

Where the real Mock import lives:

from unittest.mock import Mock

How It Works

  • We never construct a real Anthropic client and never touch the network.
  • fake_client.messages.create.return_value = fake_response tells the mock exactly what to hand back when called — no real model runs.
  • The test verifies two things: the function correctly extracts .content[0].text from the response shape, and it called the API with exactly the arguments we expect (assert_called_once_with).

🤖 How Is This Used in AI?

This is the standard, real-world way production AI codebases test anything that calls an LLM: mock the client, control exactly what it returns, and verify your own code’s logic (parsing, prompt-building, error handling) — completely independent of whether the actual model would agree with the fake response.

💡 Mocking Raw HTTP Calls (requests.post)

If your code doesn’t use an SDK client but calls requests.post directly (as shown in Module 11), you can mock the requests.post function using unittest.mock.patch:

from unittest.mock import patch, Mock
import requests

def fetch_status():
    response = requests.post("https://api.example.com/status")
    return response.json()["status"]

def test_fetch_status_mocked():
    with patch("requests.post") as mock_post:
        # Configure the mock response
        mock_response = Mock()
        mock_response.json.return_value = {"status": "operational"}
        mock_post.return_value = mock_response

        assert fetch_status() == "operational"
        mock_post.assert_called_once_with("https://api.example.com/status")

10. Patching — Replacing Real Objects During a Test

from unittest.mock import patch

def get_current_model_name():
    import os
    return os.environ.get("MODEL_NAME", "default-model")

def test_get_current_model_name_uses_env_var():
    with patch.dict("os.environ", {"MODEL_NAME": "claude-sonnet-4-6"}):
        assert get_current_model_name() == "claude-sonnet-4-6"

def test_get_current_model_name_falls_back_to_default():
    with patch.dict("os.environ", {}, clear=True):
        assert get_current_model_name() == "default-model"

🧠 Intuition: patch temporarily swaps out a real piece of the system (here, environment variables) for a controlled fake, only for the duration of the with block (recall Module 10’s context managers) — everything reverts back to normal the moment the test finishes.

🤖 How Is This Used in AI? Testing how your code behaves under different .env configurations, or temporarily replacing a real API client class with a fake one inside a specific test, without touching any actual global state permanently.


11. Testing Async Code

Recall Module 13 — testing async def functions needs a small addition:

pip install pytest-asyncio
import pytest
from unittest.mock import AsyncMock

async def call_llm_async(client, prompt):
    response = await client.messages.create(
        model="claude-sonnet-4-6",
        messages=[{"role": "user", "content": prompt}],
    )
    return response.content[0].text

@pytest.mark.asyncio
async def test_call_llm_async():
    fake_response = AsyncMock()
    fake_response.content = [AsyncMock(text="Async response here.")]

    fake_client = AsyncMock()
    fake_client.messages.create.return_value = fake_response

    result = await call_llm_async(fake_client, "test prompt")
    assert result == "Async response here."

🧠 Intuition: AsyncMock is exactly like Mock, but its methods can be awaited — necessary because a plain Mock’s methods aren’t compatible with async/await syntax.

🤖 How Is This Used in AI? Testing the concurrent tool-calling and batch-LLM-call patterns from Module 13 without actually running real concurrent network requests during your test suite.


12. Testing AI-Specific Logic

Pulling several modules together — testing a retry function without actually waiting or making real calls:

from unittest.mock import Mock

def call_with_retry(api_call, max_retries=3):
    last_error = None
    for attempt in range(max_retries):
        try:
            return api_call()
        except ConnectionError as e:
            last_error = e
    raise last_error

def test_call_with_retry_succeeds_after_failures():
    mock_api = Mock(side_effect=[
        ConnectionError("fail 1"),
        ConnectionError("fail 2"),
        "success!",
    ])

    result = call_with_retry(mock_api, max_retries=3)

    assert result == "success!"
    assert mock_api.call_count == 3

def test_call_with_retry_raises_after_max_attempts():
    import pytest
    mock_api = Mock(side_effect=ConnectionError("always fails"))

    with pytest.raises(ConnectionError):
        call_with_retry(mock_api, max_retries=3)

    assert mock_api.call_count == 3

🧠 Intuition: side_effect=[...] makes a mock return a different value (or raise a different exception) on each successive call — exactly what you need to simulate “fails twice, then succeeds,” without any real waiting or network flakiness involved. This test runs in milliseconds and is 100% reliable, every single run.


13. What NOT to Mock (or Test)

  • Don’t mock so much that your test no longer verifies anything meaningful — if you mock every single piece, you’re only testing that Python can call functions, not that your logic is correct.
  • Don’t write tests that call a real, live LLM API — they’re slow, cost money, and can fail due to the model’s natural response variation rather than an actual bug.
  • Do test your own logic thoroughly: prompt building, response parsing, validation, retry/error handling, relevance filtering — all of it deterministic, all of it mockable.
  • A small number of manual, occasional, real API calls (not part of the automated test suite) are still worth doing — to confirm your assumptions about the real API’s actual response shape haven’t drifted.

⚠️ Common Beginner Mistake: Writing a test that just checks the mock returns what you told it to return — that only proves Python works, not that your application logic is correct. Always test something your own function actually computes or decides (like the parsing, the retry count, the validation), not just the mock’s configured output.


Testing an AI Application at Several Levels

Test levelWhat it checksExample
Unit testDeterministic Python logicChunk overlap is correct
Integration testComponents work togetherRetriever can query a test index
EvaluationOutput quality over examplesAnswers remain grounded and useful
End-to-end testReal user pathUpload → retrieve → answer → citation

Unit tests alone cannot prove that an open-ended answer is good. Keep a small, versioned evaluation dataset with inputs, expected properties, and difficult edge cases. Compare quality, latency, and cost against explicit thresholds when changing the prompt, model, or retriever.

Mock the Place Your Code Looks Up

Patch the name in the module that uses it, not necessarily the library where it was originally defined. A mock proves how your code behaves under the fake response you designed; it does not prove the provider’s real API still works. Keep a smaller number of controlled integration tests for that boundary.

Avoid asserting an exact paragraph from a generative model. Assert stable properties instead: the schema is valid, required evidence is cited, unsafe tools are rejected, and a deterministic function receives the right values.

Module Summary

You can now write real automated tests with pytest, use fixtures and parametrization to keep tests clean and thorough, and — critically — mock LLM API calls so your test suite runs fast, free, and reliably, testing your own pipeline logic (prompt building, parsing, retries, validation) completely independent of any live network call.

AI Connection

Every function you wrote across Modules 4, 6, 9, and 11 — cleaning text, validating input, parsing responses, retrying failed calls — is exactly the kind of deterministic logic that deserves fast, reliable unit tests. Mocking is what makes that possible without an AI codebase’s test suite becoming slow, expensive, and flaky. Every serious production AI codebase tests this way.

Mini Practice

  1. Write a pytest test for a truncate(text, max_chars) function (recall Module 4) confirming it behaves correctly at, above, and below the limit.
  2. Write a test using pytest.raises confirming that a function raises ValueError when given an empty prompt.
  3. Create a pytest fixture returning a sample list of chat messages, and write two different tests that both use it.
  4. Write a mocked test for a call_llm(client, prompt) function, checking both the returned text and that client.messages.create was called with the correct arguments.
  5. Use Mock(side_effect=[...]) to test a retry function that fails once before succeeding, and assert it was called exactly twice.

Next: Module 16 — Data Cleaning for AI and Your First AI Data Pipeline — turning messy raw text into a clean, chunked, embedded dataset, end to end.

Author
TechByteByByte Editorial Team
Reviewed by
TechByteByByte Admin
Published
Last reviewed