Don't Mock Machine Learning Models In Unit Tests

Eugene Yan has spent time trying to apply normal software unit-testing habits to machine learning code, and he's found the fit is awkward at best. The core issue is structural. Regular software takes input data and hand-written logic and produces an expected output, which you can test with a simple assert. ML flips that equation: you feed in data and expected outputs, and what comes out the other end is the logic itself, baked into a model's weights. That means testing a model isn't about checking a function's output against a rule you wrote. It's about loading a blob of learned behavior and poking it to see if it still behaves the way you expect.

This is why Yan pushes back on the reflex to mock everything, a habit that's second nature when testing databases, APIs, or notification services. Sometimes you genuinely need to run the real model. He points out that different pretrained classifiers can label the same concept differently — Google's T5 NLI model marks factual consistency as class 1, while Meta's BART NLI model calls it class 2. Miss that kind of detail and your inference logic is silently wrong, no matter how clean your mocks look. The catch is that real models can be huge, sometimes billions of parameters, which makes loading them for every quick test painfully slow and memory-hungry.

His workaround is a layered approach. For pure logic — train/test splitting, custom distance functions, preprocessing, postprocessing, error handling — he says define tiny sample data directly inside the test file rather than pulling in CSVs or Parquet files, so tests stay self-contained and don't break when someone edits an unrelated data file. For structural checks like output shape or device placement, he shows how to build a model from just its config, with random or empty weights, using Hugging Face's transformers or the accelerate library's init_empty_weights, skipping the cost of downloading real weights entirely. And for the things that actually matter — loss decreasing during training, a model overfitting on a small sample, a classifier's confidence scores meaning what you think they mean, a model server booting and returning sane output on a batch request — he says just test against the real thing, but mark those as slow tests reserved for pre-commit or pre-merge runs rather than every save.

He also draws a line most teams blur: don't bother testing external libraries like tokenizers, optimizers, or data loaders. Assume PyTorch and Hugging Face did their job. Save the testing effort for the parts your team actually owns and could actually break.

Readers chimed in with a sharper framing of the same problem: passing unit tests built on mocks mean nothing if the system fails at integration, and probabilistic model outputs make assertions a lose-lose — tighten them and you get noisy false failures people learn to ignore, loosen them and the test stops asserting anything real. Yan doesn't claim to have solved that tension. He's treating this as an open, evolving playbook rather than a finished methodology.