Unit Test Model Functions

You've spent hours training a model, only to have it silently fail in production because a predict() function returned a list instead of a numpy array, or a preprocessing step changed shape on a new data batch. Without tests, every change to your model pipeline is a gamble. This lesson shows you how to write unit tests for model functions — not just to check outputs, but to lock down the contract between your code and your data, so future you (or your teammates) can refactor with confidence. By the end, you'll have a test suite that catches regressions early, documents expected behavior, and makes your AI engineering more robust.

The problem this lesson solves

Model functions are often treated as black boxes — you feed in data, get predictions, and hope for the best. But black boxes ship bugs. The cost of a bug in a model function is high: silent degraded predictions, wrong data types in downstream API responses, or memory blowups on large inputs. When you iterate on prompts, preprocessing, or model architecture, a small change can break assumptions elsewhere in the pipeline. Unit tests act as a safety net, catching these issues at development time instead of after deployment. They also serve as living documentation, showing exactly what your model functions should return given specific inputs. Without them, you're navigating a minefield blindfolded.

Core concept / mental model

Think of a model function as a contract: it promises to transform a given input into a specific output. Unit tests verify that contract in isolation. The mental model here is input → behavior → output. You define a representative set of inputs, call the function, and assert on the output's type, shape, value, or side effects. For AI pipelines, this means testing not just the final prediction, but each stage: preprocessing, feature extraction, the model call itself (with mocked or stubbed model), and postprocessing. The test should not depend on the actual model weights — that's an integration test. Instead, you test the logic around the model, using controlled inputs and expected results. This keeps tests fast, deterministic, and independent of network or model availability.

How it works step by step

Writing unit tests for model functions follows a repeatable pattern. Here's the breakdown:

1. Identify the function and its contract — What does it take in? What should it return? What are the edge cases? For example, a function preprocess_text(text: str) -> list[int] should return a list of token IDs.
2. Choose a test framework — pytest is the industry standard for Python. It's simple, powerful, and works well with AI projects.
3. Create a test file — Typically named test_<module>.py in the same directory or a tests/ folder.
4. Write test cases — Each test focuses on one behavior. Use assert statements. For AI functions, test: - Output shape and type - Boundary values (empty input, max length) - Determinism (same input → same output) - Error handling (invalid input raises the right exception)
5. Isolate heavy components — Mock the model itself using unittest.mock or a fixture that returns a fixed tensor. This keeps tests fast and offline.
6. Run and iterate — Execute pytest, watch it fail (TDD optional), fix code, and watch it pass. Apply the red-green-refactor cycle.

Here's a concrete flow: You have a function predict_rating(review_text). You write a test that expects predict_rating("Great movie") to return a float between 1 and 5. Then you implement the function to satisfy that test. That's the essence of test-driven AI development.

Hands-on walkthrough

Let's build a simple but realistic test suite for a model function pipeline. We'll use pytest and unittest.mock to simulate the model.

First, ensure you have pytest installed (and optionally numpy for array assertions). We'll test a tokenizer and a prediction wrapper.

pip install pytest

Create a file model_funcs.py containing the functions you want to test:

# model_funcs.py
import numpy as np

def preprocess_text(text: str, max_len: int = 10) -> list[int]:
    """Simplified tokenizer: maps characters to ASCII codes."""
    if not isinstance(text, str):
        raise TypeError("text must be a string")
    ids = [ord(c) for c in text][:max_len]
    return ids

def predict_rating(review_text: str, model=None) -> float:
    """Returns a rating between 1 and 5."""
    tokens = preprocess_text(review_text, max_len=10)
    # In real life, model() takes tokens and returns logits
    if model is None:
        # Dummy behavior for demo
        return float(sum(tokens) % 5 + 1)
    logits = model(tokens)
    score = float(np.mean(logits))
    return max(1.0, min(5.0, score))

Now create test_model_funcs.py:

# test_model_funcs.py
import pytest
from model_funcs import preprocess_text, predict_rating

class TestPreprocessText:
    def test_returns_list_of_ints(self):
        result = preprocess_text("hello")
        assert isinstance(result, list)
        assert all(isinstance(i, int) for i in result)

    def test_truncates_to_max_len(self):
        long_text = "a" * 100
        assert len(preprocess_text(long_text)) == 10

    def test_raises_on_non_string(self):
        with pytest.raises(TypeError):
            preprocess_text(123)

class TestPredictRating:
    def test_without_model_returns_float_in_range(self):
        rating = predict_rating("Great movie")
        assert isinstance(rating, float)
        assert 1.0 <= rating <= 5.0

    def test_with_mocked_model(self, mocker):
        # Mock the model to return fixed logits
        mock_model = mocker.Mock(return_value=[1.0, 2.0, 3.0])
        rating = predict_rating("Okay", model=mock_model)
        # mean of [1,2,3] = 2.0
        assert rating == pytest.approx(2.0)
        mock_model.assert_called_once()

Here we used mocker from pytest-mock — install it with pip install pytest-mock. Then run the tests:

pytest -v

Expected output (partial):

===== test session starts =====
test_model_funcs.py::TestPreprocessText::test_returns_list_of_ints PASSED
test_model_funcs.py::TestPreprocessText::test_truncates_to_max_len PASSED
test_model_funcs.py::TestPreprocessText::test_raises_on_non_string PASSED
test_model_funcs.py::TestPredictRating::test_without_model_returns_float_in_range PASSED
test_model_funcs.py::TestPredictRating::test_with_mocked_model PASSED

Notice how we mocked the model — the test runs instantly, no GPU or network needed. This is the core pattern: test the function's contract, not the model's internals.

For a more advanced scenario, test a preprocessing function that uses numpy:

# test_preprocessing.py
import numpy as np
import pytest

def normalize_vector(vec: np.ndarray) -> np.ndarray:
    norm = np.linalg.norm(vec)
    if norm == 0:
        raise ValueError("Cannot normalize a zero vector")
    return vec / norm

class TestNormalizeVector:
    def test_unit_length(self):
        v = np.array([3.0, 4.0])
        result = normalize_vector(v)
        assert np.isclose(np.linalg.norm(result), 1.0)

    def test_zero_vector_raises(self):
        with pytest.raises(ValueError):
            normalize_vector(np.array([0, 0, 0]))

This test suite would be run in CI (e.g., GitHub Actions) to automatically verify every push. The key is to make tests deterministic — mock random seeds, control randomness, and avoid any external dependencies.

Compare options / when to choose what

When writing unit tests for model functions, you have several choices for the testing framework and mocking strategy. Here's a comparison:

When to choose what:

• pytest is the default for new projects — use it unless you have legacy code.
• unittest is fine if you can't install third-party packages, but pytest is worth the dependency.
• hypothesis shines when you want to fuzz your preprocessing functions with random inputs to uncover hidden bugs.
• pytest-mock is a must for mocking models to keep tests fast and offline.

For most model functions, pytest + pytest-mock is the sweet spot. If you're doing heavy data validation, property-based testing with hypothesis is a nice addition.

Troubleshooting & edge cases

Common pitfalls when testing model functions:

• Non-deterministic outputs (e.g., dropout at inference, random sampling). Fix: set model.eval() or freeze the seed in the test fixture.
• Floating point precision — never use == on floats; use pytest.approx or np.isclose.
• Model downloads or network calls — tests become slow and flaky. Always mock the model or use a tiny local stub.
• Large test data — avoid loading full datasets in unit tests; use synthetic slices.
• Shape mismatches — assert on expected dimensions, not just element values.

Example of an edge case: a function that pads sequences to a fixed length. Test it with empty input and with overly long input.

def pad_sequence(tokens: list[int], max_len: int, pad_id: int = 0) -> list[int]:
    return tokens[:max_len] + [pad_id] * max(0, max_len - len(tokens))

class TestPadSequence:
    def test_empty_input(self):
        assert pad_sequence([], 5) == [0, 0, 0, 0, 0]
    def test_overlong_input(self):
        assert pad_sequence([1,2,3,4], 2) == [1,2]
    def test_exact_length(self):
        assert pad_sequence([1,2,3], 3) == [1,2,3]

If a test fails unexpectedly, debug by printing intermediate values or running just that test with pytest -k test_name -s. Read the traceback carefully — it often points directly to the contract violation.

What you learned & what's next

Now you can write unit tests for model functions that verify inputs, outputs, types, shapes, and error handling. You learned to mock models for fast, isolated tests, and you can apply the same pattern to any AI function — preprocessing, inference, postprocessing. You've covered the first learning objective: explain the core idea behind testing model functions. And you completed a practical exercise, meeting the second objective. Next in the Applied AI engineering track, you'll likely build on this by writing integration tests that exercise the full pipeline, or adding test coverage to a real model deployment. Whatever the next lesson, your new test suite will keep your model functions reliable and your iterations safe.

Keep your tests as part of your commit workflow, and you'll thank yourself later when a refactor breaks nothing.