Writing unit tests


Teaching: 40 min
Exercises: 60 min
  • What is a unit test?

  • How do I write and run unit tests?

  • How can I avoid test duplication and ensure isolation?

  • How can I run tests automatically and measure their coverage?

  • Implement effective unit tests using pytest

  • Execute tests in Visual Studio Code

  • Explain the issues relating to non-deterministic code

  • Implement fixtures and test parametrisation using pytest decorators

  • Configure git hooks and pytest-coverage

  • Apply best-practices when setting up a new Python project

  • Recognise analogous tools for other programming languages

  • Apply testing to Jupyter notebooks


Unit testing validates, in isolation, functionally independent components of a program.

In this lesson we’ll demonstrate how to write and execute unit tests for a simple scientific code.

This involves making some technical decisions…

Test frameworks

We’ll use pytest as our test framework. It’s powerful but also user friendly.

For comparison: testing using assert statements:

from temperature import to_fahrenheit

assert to_fahrenheit(30) == 86

Testing using the built-in unittest library:

from temperature import to_fahrenheit
import unittest

class TestTemperature(unittest.TestCase):
    def test_to_farenheit(self):
        self.assertEqual(to_fahrenheit(30), 86)

Testing using pytest:

from temperature import to_fahrenheit

def test_answer():
    assert to_fahrenheit(30) == 86

Why use a test framework?

Projects that use pytest:

Learning by example

Reading the test suites of mature projects is a good way to learn about testing methodologies and frameworks

Code editors

We’ve chosen Visual Studio Code as our editor. It’s free, open source, cross-platform and has excellent Python (and pytest) support. It also has built-in Git integration, can be used to edit files on remote systems (e.g. HPC), and handles Jupyter notebooks (plus many more formats).

Demonstration of pytest + VS Code + coverage

  • Test discovery, status indicators and ability to run tests inline
  • Code navigation (“Go to Definition”)
  • The Test perspective and Test Output
  • Maximising coverage (assert recursive_fibonacci(7) == 13)
  • Test-driven development: adding and fixing a new test (test_negative_number)

A tour of pytest

Checking for exceptions

If a function invocation is expected to throw an exception it can be wrapped with a pytest raises block:

def test_non_numeric_input():
    with pytest.raises(TypeError):


Similar test invocations can be grouped together to avoid repetition. Note how the parameters are named, and “injected” by pytest into the test function at runtime:

@pytest.mark.parametrize("number,expected", [(0, 0), (1, 1), (2, 1), (3, 2)])
def test_initial_numbers(number, expected):
    assert recursive_fibonacci(number) == expected

This corresponds to running the same test with different parameters, and is our first example of a pytest decorator (@pytest). Decorators are a special syntax used in Python to modify the behaviour of the function, without modifying the code of the function itself.

Skipping tests and ignoring failures

Sometimes it is useful to skip tests (conditionally or otherwise), or ignore failures (for example if you’re in the middle of refactoring some code).

This can be achieved using other @pytest.mark annotations e.g.

@pytest.mark.skipif(sys.platform == "win32", reason="does not run on windows")
def test_linux_only_features():

def test_unimplemented_code():


Code refactoring is “the process of restructuring existing computer code without changing its external behavior” and is often required to make code more amenable to testing.


If multiple tests require access to the same data, or a resource that is expensive to initialise, then it can be provided via a fixture. These can be cached by modifying the scope of the fixture. See this example from Devito:

def grid(shape=(11, 11)):
    return Grid(shape=shape)

def test_forward(grid):
    grid.data[0, :] = 1.

def test_backward(grid):
    grid.data[-1, :] = 7.

This corresponds to providing the same parameters to different tests.


It’s common for scientific codes to perform estimation by simulation or other means. pytest can check for approximate equality:

def test_approximate_pi():
    assert 22/7 == pytest.approx(math.pi, abs=1e-2)

Random numbers

If your simulation or approximation technique depends on random numbers then consistently seeding your generator can help with testing. See random.seed() for an example or the pytest-randomly plugin for a more comprehensive solution.


pytest has automatic integration with the Python’s standard doctest module when invoked with the --doctest-modules argument. This is a nice way to provide examples of how to use a library, via interactive examples in docstrings:

def recursive_fibonacci(n):
    """Return the nth number of the fibonacci sequence

    >>> recursive_fibonacci(7)
    return n if n <=1 else recursive_fibonacci(n - 1) + recursive_fibonacci(n - 2)

Hands-on unit testing

Getting started

Setting up our editor

  1. If you haven’t already, see the setup guide for instructions on how to install Visual Studio Code and conda.
  2. Download and extract this zip file. If using an ICT managed PC please be sure to do this in your user area on the C: drive i.e. C:\Users\your_username
    • Note that files placed here are not persistent so you must remember to take a copy before logging out
  3. In Visual Studio Code go to File > Open Folder… and find the files you just extracted.
  4. If you see an alert “This workspace has extension recommendations.” click Install All and then switch back to the Explorer perspective by clicking the top icon on the left-hand toolbar
  5. Open Anaconda Prompt (Windows), or a terminal (Mac or Linux) and run:

    conda env create --file [path to environment.yml]

    The [path to environment.yml] can be obtained by right-clicking the file name in the left pane of Visual Studio Code and choosing “Copy Path”. Right click on the command line interface to paste.

  6. Important: After the environment has been created go to View > Command Palette in VS Code, start typing “Python: Select interpreter” and hit enter. From the list select your newly created environment “diffusion”

Running the tests

  1. Open test_diffusion.py
  2. You should now be able to click on Run Test above the test_heat() function and see a warning symbol appear, indicating that the test is currently failing. You may have to wait a moment for Run Test to appear.
  3. Switch to the Test perspective by clicking on the flask icon on the left-hand toolbar. From here you can Run All Tests, and Show Test Output to view the coverage report (see Lesson 1 for more on coverage)
  4. Important: If you aren’t able to run the test then please ask a demonstrator for help. It is essential for the next exercise.

Introduction to your challenge

You have inherited some buggy code from a previous member of your research group: it has a unit test but it is currently failing. Your job is to refactor the code and write some extra tests in order to identify the problem, fix the code and make it more robust.

The code solves the heat equation, also known as the “Hello World” of Scientific Computing. It models transient heat conduction in a metal rod i.e. it describes the temperature at a distance from one end of the rod at a given time, according to some initial and boundary temperatures and the thermal diffusivity of the material:

Metal Rod

The function heat() in diffusion.py attempts to implement a step-wise numerical approximation via a finite difference method:


This relates the temperature u at a specific location i and time point t to the temperature at the previous time point and neighbouring locations. r is defined as follows, where α is the thermal diffusivity: r=\frac{\alpha
\Delta t}{\Delta

The test_heat() function in test_diffusion.py compares this approximation with the exact (analytical) solution for the boundary conditions (i.e. the temperature at ends of the end being fixed at zero). The test is correct but failing - indicating that there is a bug in the code.

Testing (and fixing!) the code

Work by yourself or with a partner on these test-driven development tasks. Don’t hesitate to ask a demonstrator if you get stuck!

Separation of concerns

First we’ll refactor the code, increasing its modularity. We’ll extract the code that performs a single time step into a new function that can be verified in isolation via a new unit test:

  1. In diffusion.py move the logic that updates u within the loop in the heat() function to a new top-level function:

    def step(u, dx, dt, alpha):

    Hint: the loop in heat() should now look like this:

    for _ in range(nt - 1):
        u = step(u, dx, dt, alpha)
  2. Run the existing test to ensure that it executes without any Python errors. It should still fail.
  3. Add a test for our new step() function:

    def test_step():
        assert step() == 

    It should call step() with suitable values for u (the temperatures at time t), dx, dt and alpha. It should assert that the resulting temperatures (i.e. at time t+1) match those suggested by the equation above. Use approx if necessary. _Hint: step([0, 1, 1, 0], 0.04, 0.02, 0.01) is a suitable invocation. It will return an array of the form [0, ?, ?, 0]. You’ll need to calculate the missing values manually using the equation in order to compare the expected and actual values.

  4. Assuming that this test fails, fix it by changing the code in the step() function to match the equation - correcting the original bug. Once you’ve done this all the tests should pass.

Solution 1

Your test should look something like this:

# test_diffusion.py
def test_step():
    assert step([0, 1, 1, 0], 0.04, 0.02, 0.01) == [0, 0.875, 0.875, 0]

Your final (fixed!) step() function should look like this. The original error was a result of some over-zealous copy-and-pasting.

# diffusion.py
def step(u, dx, dt, alpha):
    r = alpha * dt / dx ** 2

    return (
        + [
            r * u[i + 1] + (1 - 2 * r) * u[i] + r * u[i - 1]
            for i in range(1, len(u) - 1)
        + u[-1:]

Now we’ll add some further tests to ensure the code is more suitable for publication.

Testing for exceptions

We want the step() function to raise an Exception when the following stability condition isn’t met: r\leq\frac{1}{2} Add a new test test_step_unstable, similar to test_step but that invokes step with an alpha equal to 0.1 and expects an Exception to be raised. Check that this test fails before making it pass by modifying diffusion.py to raise an Exception appropriately.

Solution 2

# test_diffusion.py
def test_step_unstable():
    with pytest.raises(Exception):
        step([0, 1, 1, 0], 0.04, 0.02, 0.1)

# diffusion.py
def step(u, dx, dt, alpha):
    r = alpha * dt / dx ** 2

    if r > 0.5:
        raise Exception


Adding parametrisation

Parametrise test_heat() to ensure the approximation is valid for some other combinations of L and tmax (ensuring that the stability condition remains true).

Solution 3

# test_diffusion.py
@pytest.mark.parametrize("L,tmax", [(1, 0.5), (2, 0.5), (1, 1)])
def test_heat(L, tmax):
    nt = 10
    nx = 20
    alpha = 0.01


After completing these two steps check the coverage of your tests via the Test Output panel - it should be 100%.

The full, final versions of diffusion.py and test_diffusion.py are available on GitHub.

Bonus task(s)

  • Write a doctest-compatible docstring for step() or heat()
  • Write at least one test for our currently untested linspace() function

Advanced topics

More pytest plugins

def test_fibonacci(benchmark):
    result = benchmark(fibonacci, 7)
    assert result == 13

pytest-benchmark example

Demonstration of performance regression via recursive and formulaic approaches to Fibonacci calculation (output)

from fibonacci import recursive_fibonacci
from hypothesis import given, strategies

def test_recursive_fibonacci(n):
    phi = (5 ** 0.5 + 1) / 2
    assert recursive_fibonacci(n) == int((phi ** n - -phi ** -n) / 5 ** 0.5)

Taking testing further

Testing in other languages

Further resources

Key Points

  • Testing is not only standard practice in mainstream software engineering, it also provides distinct benefits for any non-trivial research software

  • pytest is a powerful testing framework, with more functionality than Python’s built-in features while still catering for simple use cases

  • Testing with VS Code is straightforward and encourages good habits (writing tests as you code, and simplifying test-driven development)

  • It is important to have a set-up you can use for every project - so that it becomes as routine in your workflow as version control itself

  • pytest has a myriad of extensions that are worthy of mention such as Jupyter, benchmark, Hypothesis etc

  • Adding unit tests can verify the correctness of software and improve its structure: isolating logical distinct code for testing often involves untangling complex structures