Jacob Kaplan-Moss

Getting started with pytest

I wrote this post in 2016, more than 7 years ago. It may be very out of date, partially or totally incorrect. I may even no longer agree with this, or might approach things differently if I wrote this post today. I rarely edit posts after writing them, but if I have there'll be a note at the bottom about what I changed and why. If something in this post is actively harmful or dangerous please get in touch and I'll fix it.

Pytest is my preferred Python testing library. It makes simple tests incredibly easy to write, and is full of advanced features (and tons of plugins) that help with more advanced testing scenarios.

To demonstrate the basics, I’m going to walk through how I’d solve the first couple cryptopals challenges in a test-driven style, using py.test.

Spoiler alert:

I’m going to spoil the first challenge, and maybe a bit of the second, below. If you want to work through them yourself, do that before reading the rest of this post.

Installation and a first test

Installation is typical:

$ pip install pytest

Note

I’m using Python 3.5, and I’m doing all this in a virtualenv. If you don’t have Python installed, or don’t already know how to use pip and virtualenv, check out The Hitchiker’s Guide to Python for a good installation guide.

The first challenge asks us to convert a hex-encoded string to base64. I’ll start by writing a test to represent the challenge. py.test by default looks for tests in a files named something like test_whatever.py, so I’ll make a test_set1.py and write my test there:

import base64

def test_challenge1():
    given = "49276d206b696c6c696e6720796f757220627261696e206c696b65206120706f69736f6e6f7573206d757368726f6f6d"
    expected = b"SSdtIGtpbGxpbmcgeW91ciBicmFpbiBsaWtlIGEgcG9pc29ub3VzIG11c2hyb29t"
    assert base64.b64encode(bytes.fromhex(given)) == expected

Did I get it right?

$ py.test
=============================================== test session starts ================================================
platform darwin -- Python 3.5.2, pytest-3.0.4, py-1.4.31, pluggy-0.4.0
rootdir: /Users/jacobkaplan-moss/c/pytest-blog, inifile:
collected 1 items

test_set1.py .

Yes, I haven’t really written any code yet: the first challenge is super-simple in Python, which can parse a hex-encoded string to a bytestring using bytes.fromhex, and has base64 encoding build-in as the base64 module.

However, this demonstrates the “simple” part of pytest: tests are just simple functions named test_whatever(), and rather than having than a bunch of assert methods (assertEqual, assertNotEqual, assertAlmostEqual), you just write simple assert statements.

A second, more realistic testing situation

To see a more complete example, I’ll solve the second challenge, which asks to implement a function that XORs two fixed-length buffers. In test-driven style, I’ll write the test first:

from cryptopals import fixed_xor

def test_challenge2():
    bs1 = bytes.fromhex("1c0111001f010100061a024b53535009181c")
    bs2 = bytes.fromhex("686974207468652062756c6c277320657965")
    assert fixed_xor(bs1, bs2).hex() == "746865206b696420646f6e277420706c6179"

In typical test-driven style, I’ll now immediately run tests:

$ py.test

====================================================== ERRORS ======================================================
__________________________________________ ERROR collecting test_set1.py ___________________________________________
ImportError while importing test module '/Users/jacobkaplan-moss/c/pytest-blog/test_set1.py'.
Hint: make sure your test modules/packages have valid Python names.
Traceback:
test_set1.py:2: in <module>
    from cryptopals import fixed_xor
E   ImportError: No module named 'cryptopals'
!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!! Interrupted: 1 errors during collection !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
1 error in 0.12 seconds

As expected, I get an error – I haven’t created the cryptopals module the test tries to import. This error looks different from a failed test because this is happening during what pytest calls the “collection” phase. This is when pytest walks through your files, looking for test modules (files named test_whatever.py) and test functions (def test_whatever()). For more on how pytest discovers tests (and how to customize it), see conventions for Python test discovery.

I’ll now stub out this function, mostly to demonstrate what test failure looks like. In cryptopals.py, I wrote:

def fixed_xor(bs1, bs2):
    return b''

And, as expected, I get a failure:

$ py.test -q
.F
===================================================== FAILURES =====================================================
_________________________________________________ test_challenge2 __________________________________________________

    def test_challenge2():
        bs1 = bytes.fromhex("1c0111001f010100061a024b53535009181c")
        bs2 = bytes.fromhex("686974207468652062756c6c277320657965")
>       assert fixed_xor(bs1, bs2).hex() == "746865206b696420646f6e277420706c6179"
E       assert '' == '746865206b696420646f6e277420706c6179'
E         + 746865206b696420646f6e277420706c6179

test_set1.py:12: AssertionError
1 failed, 1 passed in 0.04 seconds

(I’m using -q – short for --quiet – to get slightly less output.)

I love the way that this highlights the line where the test failed on, and shows me the values of what that assert statement ran on. pytest is doing some tremendously dark black magic to make this happen, and the result is super-great.

Once I write the correct code (omitted here in the spirit of keeping these challenges challenging), I should see the following:

$ py.test -q
..
2 passed in 0.01 seconds

A taste of more advanced pytest: parameterized test functions

For a final example, as a way of looking at a slightly more complex use of pytest, I want to write a few more test to check what happens when my fixed_xor function gets fed bytestrings of different lengths. The challenge only says that the function should “takes two equal-length buffers”, but doesn’t specify what happens when those buffers aren’t the same length. So, I made the decision that the result should be the length of the shortest bytestring (mostly because that makes the function easier to write).

To test this properly, I should test against a few different scenarios: bs1 being shorter than bs2, len(bs2) < len(bs1), and where either bs1 or bs2 are empty – corner cases are always where bugs lurk! I could write four more test functions, but that’s repetitive. So I’ll turn to parameterized test functions:

import pytest

@pytest.mark.parametrize("in1, in2, expected", [
    ("1c011100", "686974207468652062756c6c277320657965", "74686520"),
    ("1c0111001f010100061a024b53535009181c", "68697420", "74686520"),
    ("", "68697420", ""),
    ("1c011100", "", ""),
    ("", "", "")
])
def test_challenge2_mismatching_lengths(in1, in2, expected):
    bs1 = bytes.fromhex(in1)
    bs2 = bytes.fromhex(in2)
    assert fixed_xor(bs1, bs2) == bytes.fromhex(expected)

This is the general pattern for doing more advanced work in pytest: use a decorator to somehow annotate or modify the test function to do something special. Here, the parametrize decorator lets me specify a list of arguments to be passed to the test function; the test function will then be run once for each set of parameters. Notice what happens when I run the tests:

$ py.test -q
.......
7 passed in 0.01 seconds

Rather than just showing test_challenge2_mismatching_lengths as a single test, we see five tests – one for each example. Because each set of parameters shows up as a separate case, if I add another example designed to deliberately fail, I’ll just see that one failure, and know exactly what it was:

$ py.test -q
.......F
===================================================== FAILURES =====================================================
_________________________ test_challenge2_mismatching_lengths[1c011100-68697420-12345678] __________________________

in1 = '1c011100', in2 = '68697420', expected = '12345678'

    @pytest.mark.parametrize("in1, in2, expected", [
        ("1c011100", "686974207468652062756c6c277320657965", "74686520"),
        ("", "68697420", ""),
        ("", "68697420", ""),
        ("1c011100", "", ""),
        ("", "", ""),
        ("1c011100", "68697420", "12345678"),
    ])
    def test_challenge2_mismatching_lengths(in1, in2, expected):
        bs1 = bytes.fromhex(in1)
        bs2 = bytes.fromhex(in2)
>       assert fixed_xor(bs1, bs2) == bytes.fromhex(expected)
E       assert b'the ' == b'\x124Vx'
E         At index 0 diff: 116 != 18
E         Full diff:
E         - b'the '
E         + b'\x124Vx'

test_set1.py:26: AssertionError
1 failed, 7 passed in 0.05 seconds

pytest is full of niceties like this – different ways to easily manage setup/teardown scenarios, ways to share resources between different test modules, tons of options on how to organize and factor test code, ways to group and mark tests, and so on. It’s a great library that really makes writing test code easy and pleasant. I hope you’ll check it out!