Better testing through concentration

Basically every piece of software I write eventually uses some form of stochasticity. Writing tests of random functions is frustrating: you write a test that draws 1,000 samples¹ and asserts it’s the right value (or within some tolerance I pick by feel). It passes and I move on with my life. Then, three months later, when I need to push a small fix somewhere, it starts failing and I have no idea why.

The usual response is either (a) crank up the sample count until the failures stop, (b) widen the tolerance until the test is effectively vacuous, or (c) mark the test as flaky and move on with my life². I’m sure this sounds familiar, but these solutions suck. They’re fundamentally symptoms of the same underlying problem: there’s no principled connection between the test’s statistical properties and the sample count and tolerance chosen.

This has always been annoying, but it’s become a more urgent problem recently. When writing code with AI, whether code is actually good tracks the quality of tests and evaluations. Strong tests mean strong code: the model writes something, the tests catch what’s wrong, the model fixes it, and you converge to a working piece of software at the end of the session. Bad, vacuous tests break this loop. The model “passes” the tests by writing a test that always passes (“Oh, this test is flaky. I’ll just crank up the tolerance.” Yes, unfortunately, Claude has the same bad habits I do, sometimes), and you don’t find out until production. Stochastic code is especially vulnerable here, because the failure mode is silent: a test with a too-wide tolerance or a too-small sample count may happily pass code that’s subtly wrong, but you’ll feel confident because you have a test written supposedly verifying its behavior.

This bothered me enough to (try to) do something about it.³

The idea

pytest-stochastic is a pytest plugin that makes the connection between test and statistical guarantee explicit. You declare what you know about your test statistic—its bounds, variance, sub-Gaussian parameter, whatever you’ve got—and the framework does the rest. It selects the tightest applicable concentration inequality from a registry of ten I chose “stochastically”, computes the minimum sample size to hit the acceptable flakiness budget, and runs the test.

from pytest_stochastic import stochastic_test

@stochastic_test(
    expected=0.5,
    atol=0.01,
    bounds=(0, 1),
    failure_prob=1e-8,
)
def test_uniform_mean(rng):
    return rng.uniform(0, 1)

That failure_prob=1e-8 is guaranteed based on the invariants you provide: the test will produce a false failure less than once in 100 million runs (so long as your assumptions are correct). The framework figures out that it needs exactly 7,394 samples to deliver that guarantee via Hoeffding’s inequality, and runs them.

The key insight is that you don’t need to carefully figure out the right concentration inequality, you just express the assumptions that you know. If you only know the range of your statistic, you get Hoeffding and pay for the worst case. Declare the variance and Bernstein kicks in—often much tighter. Know the sub-Gaussian parameter and you can do even better still. In the worked example in the docs, testing an OLS slope estimator goes from 75,000 samples (bounds only) to 100 (sub-Gaussian): a reduction of multiple orders of magnitude for the same probabilistic guarantee.

But I don’t know my variance

Yeah, no shit, why would you? Deriving variance analytically for a unit test is absurdly overkill. This is what tune mode is for:

pytest --stochastic-tune

This profiles each test empirically (50,000 samples by default), computes an upper confidence bound on the variance, and writes the result to .stochastic.toml. Commit that file, and subsequent runs automatically use the tighter Bernstein bound.

What else

There’s also a @distributional_test decorator for goodness-of-fit testing (KS, chi-squared, Anderson-Darling against a reference distribution), automatic RNG injection with seed reporting on failure for reproducibility, and sensible error messages when you misconfigure things. The documentation covers all of this.

Why bother

Ad-hoc thresholds and magic sample counts are technical debt that compounds quietly, and this is an attempt to do better. For many use cases, alas, the sample counts that concentration inequalities require are still too large to be practical. I tried this on a recent project and ended up just fixing a seed anyway, because I couldn’t afford 10–100 evaluations of the test statistic (which each required letting a neural network converge). When individual evaluations are cheap, though, this should be helpful to get faster, more reliable tests (or at least figure out what the tradeoff is between the two). I think the worked example in the docs is a pretty good example of how this could be useful.

This isn’t new math or anything, but maybe someone finds it useful.

The code is on GitHub. Try it out and let me know about how it goes. I’m making no guarantees, here!

Footnotes

1. Why 1,000? Because it’s a nice round number, and nice round numbers are inherently trustworthy.

2. Or more likely, just choose a random seed that works for a single run and hardcode it. It’s future-Drew’s problem when a change makes it randomly start failing.

3. Read as: bothered me enough to make Claude do something about it. Pretty nearly one-shot the thing, too, but as always the long-tail of iteration still took a while. GitHub shows 15 PRs and about 50 commits, so that’s probably about how much prompting I had to do (each PR represented one initial prompt that I iterated on a couple times fixing errors and providing feedback).