Testing TensorFlow Code and Models: Unit Tests, Golden Files, and CI That Catches Regressions

Every TensorFlow team we work with has a test suite for the web app in front of the model and nothing at all for the model itself. Then a refactor of a preprocessing layer silently changes tokenization, the offline metric looks fine because it was computed with the same broken code, and three weeks later someone notices conversions fell off a cliff. Machine learning code fails quietly. Testing is how you make it fail loudly, in CI, before a reviewer has to trust a screenshot of a loss curve.

This tutorial builds a practical test pyramid for a Keras 3 / TensorFlow project: fast unit tests for layers and pipelines, golden-file tests for serving signatures, behavioural tests for model quality, and a GitHub Actions workflow that runs the cheap ones on every push and the expensive ones nightly.

What to test (and what not to)

You are not testing that gradient descent works. You are testing the code around the math, which is where the bugs are:

Layer of the pyramid	Runs in	Catches
Unit tests: custom layers, losses, preprocessing fns	seconds, CPU	shape/dtype bugs, broken masking, non-serializable layers
Data tests: schema, ranges, leakage checks	seconds, CPU	upstream schema drift, label leakage
Contract tests: SavedModel signature, golden outputs	~a minute, CPU	silent input-format changes, preprocessing drift
Training smoke test: 1 step, tiny data	~a minute, CPU/GPU	broken fit, bad shapes at scale, NaNs
Behavioural / quality gates: metrics on a frozen slice	minutes, GPU	accuracy regressions, fairness slice regressions

Anything above the line runs on every pull request. Quality gates run nightly or on release branches, because they need real data and a GPU.

1. Unit-test custom layers

A custom layer has three contracts: output shape, dtype behaviour, and round-trip serialization. Test all three. The third one is the one everybody forgets, and it is the one that breaks model.save() at 6pm on a Friday.

import numpy as np
import keras
import tensorflow as tf

from myproject.layers import GatedResidual


class GatedResidualTest(tf.test.TestCase):

    def test_output_shape_and_finiteness(self):
        layer = GatedResidual(units=16)
        x = keras.random.normal((4, 10, 32))
        y = layer(x)
        self.assertEqual(y.shape, (4, 10, 16))
        self.assertAllInRange(y, -1e4, 1e4)   # no NaN/Inf sneaking through

    def test_serialization_round_trip(self):
        layer = GatedResidual(units=16)
        cfg = layer.get_config()
        clone = GatedResidual.from_config(cfg)
        self.assertEqual(clone.units, layer.units)

    def test_masking_is_respected(self):
        layer = GatedResidual(units=8)
        x = keras.random.normal((2, 5, 4))
        mask = np.array([[True] * 3 + [False] * 2] * 2)
        y = layer(x, mask=mask)
        # padded steps must not influence the pooled summary
        x2 = np.array(x)
        x2[:, 3:, :] = 999.0
        y2 = layer(x2, mask=mask)
        self.assertAllClose(y, y2, atol=1e-5)

tf.test.TestCase gives you assertAllClose, assertAllInRange, and deterministic seeding helpers; it is a unittest.TestCase, so pytest collects it happily. Keep these tests on CPU and keep them under a second each — a suite you have to wait for is a suite people stop running.

Two more unit tests worth writing for every project:

• Loss symmetry / edge cases. Feed a perfect prediction and assert the loss is ~0; feed the worst case and assert it is finite, not inf.
• Preprocessing determinism. Call your tf.data map function twice on the same record and assert identical output (unless it is deliberately random, in which case assert it is not identical under different seeds and is identical under the same seed).

2. Test the data, not just the code

Schema and range checks belong in CI against a small committed sample, and in the pipeline against the real feed. If you already run TensorFlow Data Validation for drift — see our tutorial on catching data drift early with TFDV — reuse the same schema proto as the CI fixture so the two can never disagree:

import tensorflow_data_validation as tfdv

def test_sample_matches_schema():
    stats = tfdv.generate_statistics_from_csv("tests/fixtures/sample.csv")
    schema = tfdv.load_schema_text("schema/train_schema.pbtxt")
    anomalies = tfdv.validate_statistics(stats, schema)
    assert not anomalies.anomaly_info, tfdv.utils.display_util.get_anomalies_dataframe(anomalies)

Add one leakage test that is specific to your problem: assert that no feature column correlates with the label above some absurd threshold, or that the timestamp of every feature precedes the label timestamp. Leakage is the single most expensive ML bug we get called in to diagnose, and a five-line test catches most instances of it.

3. Golden tests for the serving signature

The SavedModel signature is your API contract with everything downstream. Freeze it. Commit a small fixture of inputs and the expected outputs, and fail the build if either the signature keys or the numbers move.

import json, pathlib
import numpy as np
import tensorflow as tf

GOLDEN = pathlib.Path("tests/golden/predictions.json")

def test_signature_contract(tmp_path):
    model = build_and_load_weights()           # your project's factory
    export_dir = tmp_path / "1"
    model.export(export_dir)                   # Keras 3 -> SavedModel

    loaded = tf.saved_model.load(str(export_dir))
    fn = loaded.signatures["serving_default"]

    # 1. the contract itself
    assert set(fn.structured_input_signature[1]) == {"tokens", "numeric"}
    assert set(fn.structured_outputs) == {"score"}

    # 2. the numbers
    fixture = json.loads(GOLDEN.read_text())
    out = fn(
        tokens=tf.constant(fixture["tokens"]),
        numeric=tf.constant(fixture["numeric"], dtype=tf.float32),
    )["score"].numpy()
    np.testing.assert_allclose(out, fixture["score"], rtol=1e-4, atol=1e-5)

Regenerate the golden file deliberately, with a script, in its own commit, so the diff shows a reviewer exactly how much the predictions moved. A golden file that gets silently rewritten by the test run is worse than no test.

Tolerances matter: exact equality will flake across GPU kernels, XLA versions, and mixed precision. rtol=1e-4 is a reasonable default for float32; loosen to rtol=1e-2 for bfloat16 paths, and if you are testing a quantized LiteRT export, compare decisions (argmax, thresholded label) rather than raw floats.

4. A one-step training smoke test

The cheapest high-value test in the suite: run the real training entry point for one step on two batches of synthetic data.

def test_train_step_runs_and_changes_weights():
    model = build_model()
    ds = synthetic_dataset(batch_size=2, steps=2)
    before = [w.numpy().copy() for w in model.trainable_weights]

    history = model.fit(ds, epochs=1, steps_per_epoch=2, verbose=0)

    loss = history.history["loss"][-1]
    assert np.isfinite(loss)
    after = [w.numpy() for w in model.trainable_weights]
    assert any(not np.allclose(b, a) for b, a in zip(before, after)), "no weights updated"

That last assertion catches the classic frozen-layer mistake, where a fine-tune trains nothing because everything was left non-trainable. Add a variant with jit_compile=True if you rely on XLA, because XLA compilation failures only show up at step one.

5. Quality gates with real thresholds

Behavioural tests answer "is this model good enough to ship", and they need real data, so they run nightly or on a release branch, not on every push. Three rules make them useful instead of annoying:

1. Gate on a frozen evaluation slice, stored with a content hash. If the eval set changes, the comparison is meaningless.
2. Gate against the current production model, not an absolute number. new_auc >= prod_auc - 0.002 is a shippable rule; auc > 0.85 becomes a lie the moment the data shifts.
3. Gate per slice. An aggregate metric that holds while performance on your highest-value customer segment collapses is exactly the failure mode TFMA's sliced metrics exist to catch.

THRESHOLDS = {"overall": 0.002, "segment:enterprise": 0.005}

def test_no_metric_regression(candidate_metrics, production_metrics):
    for slice_name, allowed_drop in THRESHOLDS.items():
        delta = production_metrics[slice_name] - candidate_metrics[slice_name]
        assert delta <= allowed_drop, f"{slice_name} regressed by {delta:.4f}"

In a TFX pipeline this is the Evaluator component with a tfma.MetricThreshold and change_threshold, and the blessing artifact it produces is what gates the Pusher. If you are not on TFX yet, the same logic as a pytest job against a metrics JSON in object storage gets you 80% of the value — see our TFX continuous training walkthrough for the full pipeline version.

6. Wiring it into CI

Split the workflow by cost. Fast tests on every push, on CPU runners; expensive tests on a schedule, on a GPU runner.

name: ml-ci
on:
  push:
  schedule:
    - cron: "0 3 * * *"

jobs:
  fast:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v5
        with:
          python-version: "3.12"
          cache: pip
      - run: pip install -r requirements-ci.txt   # CPU TensorFlow, pinned
      - run: pytest tests/unit tests/data tests/contract -q --durations=10

  nightly-quality:
    if: github.event_name == 'schedule'
    runs-on: [self-hosted, gpu]
    steps:
      - uses: actions/checkout@v4
      - run: pip install -r requirements-gpu.txt
      - run: pytest tests/quality -q -m "not flaky"
      - run: python scripts/publish_metrics.py --run-id ${{ github.run_id }}

Details that decide whether this stays green:

• Pin everything. TensorFlow, Keras, and CUDA versions in a lockfile. Keras 3 model code is portable across backends, but numerical output is not bit-identical across them, and your golden tests will notice.
• Set determinism where you can. keras.utils.set_random_seed(0) plus tf.config.experimental.enable_op_determinism() in a pytest fixture. Determinism costs speed, so enable it for the test session only.
• Silence the GPU in fast tests. CUDA_VISIBLE_DEVICES="" on the CPU job so a stray allocation cannot make the fast suite flaky.
• Mark and quarantine. A @pytest.mark.flaky test that nobody fixes is a broken window; give flaky tests an owner and a deadline, or delete them.
• Cache the fixtures, not the models. Small committed fixtures beat downloading a dataset in CI.

A reasonable starting point

If you have zero tests today, add these five, in this order, and you will have caught most of what actually breaks:

1. Serialization round-trip for every custom layer.
2. Schema validation against a committed data sample.
3. Signature + golden-output contract test on the exported SavedModel.
4. One-step training smoke test.
5. Nightly metric-regression gate against the production model, per slice.

That is a day of work for one engineer and it converts "we think the model is fine" into a check a reviewer can read. It also makes every subsequent change cheaper, which is the real argument: teams with model tests ship retrains weekly; teams without them ship quarterly and hold their breath.

We do this as a fixed-scope engagement — audit the repo, stand up the pyramid, wire the CI, hand it back with documentation. If your TensorFlow project has more model versions than tests, get in touch and we will scope it.