Pytest for AI Code Coverage

You’ve built an AI pipeline that returns beautiful answers — until a refactor silently breaks the tokenizer, or a prompt tweak halves accuracy on edge cases. You run the tests, but they only cover the happy path. That’s why use pytest for AI code coverage is not a luxury; it’s the guardrail that catches the “invisible” failures — the surprising ones, the slow ones, the ones that only appear when your model gets a slightly different input. By the end of this lesson, you’ll write a pytest suite with coverage measurement tuned specifically for AI code — and you’ll know exactly which lines matter most when the model behaves unpredictably.

The Problem This Lesson Solves

AI code is deceptive: it looks simple (a few API calls, a couple of prompts), but its failure modes are exotic. A unit test that asserts response["text"] isn’t empty passes even when your code never calls the real model — or worse, when it calls it with a malformed prompt and gets a hallucinated answer that happens to have text. The coverage report doesn’t lie: it shows you which parts of your pipeline are actually executed in tests. When you use pytest for AI code coverage, you shift from “tests pass” to “I know what my tests actually touch.” That distinction is critical because AI code has three layers of uncertainty: the model’s nondeterminism, the prompt’s fragility, and your parsing logic’s assumption about output structure. Standard coverage tools treat them all the same — but you need to see which of those layers you’re testing and which you’re blind to.

Traditional coverage questions like “did we run this if branch?” become AI-specific: “Did we test what happens when the model returns an empty JSON? When it returns extra fields? When it times out?” These are the real bugs that kill AI apps in production. This lesson gives you the practical toolkit: a pytest setup with coverage.py or pytest-cov, plus a strategy to measure and improve coverage for your AI modules — without drowning in false confidence.

Core Concept / Mental Model

Think of your AI code as a pipeline with three gates. First is the prompt gate — the template building, tokenization, and input validation. Second is the model call gate — the API request, retry logic, and timeout handling. Third is the output gate — response parsing, schema validation, and fallback logic. Coverage is a map of which gates you’ve actually walked through in tests.

A naive assert model_response("hello") might walk all three gates once. But coverage percent shows you segments of each gate. For example, the error-handling branch of your retry loop might remain untouched, so if the API rate-limits you, your test suite never exercises that path. With pytest --cov, you see a per-file line table — you can spot that the except clause on line 42 is red (uncovered). That visual is your mental model: green lines are gates you’ve opened; red lines are doors you haven’t even knocked on.

Coverage alone doesn’t guarantee correctness — an AI model never gives the same answer twice, so a passing test doesn’t mean your assertions are strong. But coverage gives you a sense of control: you know exactly which lines of your code are depended upon by tests and which are dead weight or untested risk.

How It Works Step by Step

Let’s walk through the process of setting up and using pytest with coverage for an AI module.

Step 1: Install the tools

You need pytest and pytest-cov (which wraps coverage.py for pytest). For AI projects, you also want a mockable HTTP client library, like responses or pytest-mock.

pip install pytest pytest-cov

Step 2: Write a basic pytest file

Create a simple AI pipeline — for example, a function that calls an OpenAI‑style API and extracts a JSON response.

# ai_pipeline.py
import json

def call_model(prompt: str, api_key: str):
    """Simulated model call. In production, replace with requests.post(...).
    Returns a dict with 'choices' and 'error' keys."""
    if not prompt.strip():
        raise ValueError("Prompt cannot be empty")
    if not api_key:
        raise ValueError("API key is missing")
    # Simulated response — in real code, this would be an API response
    return {"choices": [{"message": {"content": '{"weather": "sunny"}'}}]}

def parse_weather(content: str) -> dict:
    """Parse the model's content string into a dict."""
    try:
        return json.loads(content)
    except json.JSONDecodeError:
        return {"error": "Invalid JSON"}

def get_weather(prompt: str, api_key: str) -> dict:
    """End-to-end weather extraction."""
    response = call_model(prompt, api_key)
    content = response["choices"][0]["message"]["content"]
    return parse_weather(content)

Now write a test that covers the happy path and one error path.

# test_ai_pipeline.py
import pytest
from ai_pipeline import get_weather, call_model, parse_weather

def test_get_weather_happy_path():
    result = get_weather("What's the weather?", "test-key")
    assert result == {"weather": "sunny"}

def test_get_weather_empty_prompt():
    with pytest.raises(ValueError):
        get_weather("", "test-key")

def test_parse_invalid_json():
    assert parse_weather("not json") == {"error": "Invalid JSON"}

Step 3: Run pytest with coverage

Run pytest --cov=ai_pipeline to see the report for that module. You’ll see something like:

Name             Stmts   Miss  Cover
------------------------------------
ai_pipeline          12      3    75%

Now run pytest --cov=ai_pipeline --cov-report=term-missing to see which lines are missing:

Name             Stmts   Miss  Cover   Missing
----------------------------------------------
ai_pipeline          12      3    75%   12-14

Line 12–14 is likely the except block in parse_weather or the valueerror for missing API key. That tells you exactly which scenario you haven’t tested.

Step 4: Improve coverage with targeted tests

Add a test that doesn’t pass an API key to cover the missing branch. Then re-run coverage to confirm it’s now 100% for that module.

def test_missing_api_key():
    with pytest.raises(ValueError):
        get_weather("Hello", "")

Pro tip: Aim for 100% on the deterministic parts of your AI code — parsing, validation, error handling. The model call itself is external and inherently nondeterministic, so you mock it via pytest-mock or a real integration test (covered separately).

Hands-on Walkthrough

Now you’ll build a complete, more realistic example — an AI service that extracts entities from text. You’ll mock the API to control responses and use coverage to find a hidden branch that could break your app.

Project setup

Create a virtual environment and install dependencies.

python -m venv venv && source venv/bin/activate
pip install pytest pytest-cov requests

Code: entity_extractor.py

# entity_extractor.py
import requests
import json

API_URL = "https://api.example.com/extract"

def extract_entities(text: str, api_key: str) -> dict:
    """Extract named entities from text using a mock API."""
    if not text:
        raise ValueError("Text cannot be empty")
    headers = {"Authorization": f"Bearer {api_key}"}
    payload = {"text": text}
    try:
        response = requests.post(API_URL, json=payload, headers=headers, timeout=5)
        response.raise_for_status()
    except requests.exceptions.Timeout:
        return {"error": "timeout"}
    except requests.exceptions.HTTPError as e:
        return {"error": f"http_{e.response.status_code}"}
    data = response.json()
    # Assume API returns {"entities": ["OpenAI"]}
    return data.get("entities", [])

Tests with mocking

You’ll use pytest-mock to simulate different API responses, including a timeout and a 500 error, to exercise the except branches.

# test_entity_extractor.py
import pytest
from entity_extractor import extract_entities

def test_extract_success(mocker):
    mock_response = mocker.Mock()
    mock_response.raise_for_status.return_value = None
    mock_response.json.return_value = {"entities": ["OpenAI"]}
    mocker.patch("entity_extractor.requests.post", return_value=mock_response)
    result = extract_entities("OpenAI released GPT-4", "key")
    assert result == ["OpenAI"]

def test_timeout(mocker):
    mocker.patch("entity_extractor.requests.post", side_effect=requests.exceptions.Timeout)
    result = extract_entities("Anything", "key")
    assert result == {"error": "timeout"}

def test_http_error(mocker):
    mock_response = mocker.Mock()
    mock_response.raise_for_status.side_effect = requests.exceptions.HTTPError(response=mocker.Mock(status_code=500))
    mocker.patch("entity_extractor.requests.post", return_value=mock_response)
    result = extract_entities("Anything", "key")
    assert result == {"error": "http_500"}

def test_empty_text():
    with pytest.raises(ValueError):
        extract_entities("", "key")

Run coverage:

pytest --cov=entity_extractor --cov-report=term-missing

You’ll likely see 100% coverage for that module — but that doesn’t mean tested the right things. Notice the response.json() call — the success test mocked it as return_value, so the real JSON decoding path is uncovered. That’s where real bugs hide (e.g., the API returns a string instead of a dict). Add a test that uses a real requests.Response object with .json() that raises json.JSONDecodeError.

Expected output after fixing

Name               Stmts   Miss  Cover   Missing
------------------------------------------------
entity_extractor       20      0   100%

Pro tip: Use pytest --cov=entity_extractor --cov-report=html to generate an interactive HTML report — click any line to see exactly how many times it ran with which arguments. That’s invaluable when debugging a tricky AI pipeline.

Compare Options / When to Choose What

You have several ways to generate coverage reports. Here’s a quick comparison:

When to choose what: If you’re starting, use pytest-cov with --cov=your_module. If you have complex AI logic (schema validation, retry loops), add --cov-branch to catch missed if/else paths. For CI, export an XML report (--cov-report=xml) and integrate with Codecov or SonarQube.

Troubleshooting & Edge Cases

Coverage report shows 0% even though tests run

Often happens when the module is imported from a different path. Fix: ensure your test and module are in the same package, or specify the path correctly.

Coverage drops after adding a mock

That’s expected — mocking can bypass real code. Always test both the mocked and unmocked paths (e.g., use pytest-mock for one test, and a real integration test with a timeout for another).

--cov-branch gives lower percentage

Branch coverage is stricter — it counts each if branch as a separate line. It’s normal to see 60–80% on AI code. Focus on critical branches: parsing errors, missing API keys, empty responses.

AI model returns nondeterministic content

Don’t assert exact strings. Instead, use json.loads on the content or check that the response has expected keys. Coverage helps you find which lines are never executed — a hint that the model’s output shape might be different than you expected

coverage not installed

Run pip install coverage pytest-cov. If you’re in a container, remember to reinstall after rebuilding.

What You Learned & What's Next

You now know how to use pytest for AI code coverage: you can set up pytest-cov, run line‑level coverage, identify missing branches with --cov-report=term-missing, and use mocking to exercise error paths without paying for real API calls. You also learned that coverage is a map of your pipeline’s gates — prompt, model call, and output — and you should aim for 100% on the deterministic parts while treating the model call as an external dependency you control via mocks.

Next in the track: you’ll take this further with mocking AI APIs with responses — you’ll replace real HTTP calls with scripted responses to test your pipeline’s resilience to timeouts, rate limits, and malformed JSON. That gives you deterministic tests even when the model itself is black‑box.

Key takeaways to remember:

• Coverage tells you what you tested, not how well — combine it with strong assertions on output structure.
• Use pytest --cov=your_module and read the “Missing” column — that’s your to‑do list.
• Mock external AI calls to test error paths; never rely on live calls in unit tests.
• Branch coverage (--cov-branch) is essential for AI code with fallback logic.
• Keep an eye on the response.json() line — that’s where real‑world parsing bugs hide.

Practice recap

Take your own AI module (or the entity_extractor above) and run pytest --cov in its simplest form. Then add a test that covers the missing branch from the report — for example, a malformed JSON response. Re‑run coverage and watch the percentage climb. That’s the loop: run, find red lines, test them, repeat. It feels like a superpower after ten minutes.