Sam's Log

Pytest Tricks: Freezegun and Parametrize

I always try to work following Test Driven Development. I recently used Pytest to write some unit tests and discovered a couple of neat tricks from a work colleague.

Context

I needed to write a function which determined if a user’s account had been created within a specified time window. The function returns a boolean i.e. true if the account was created within the time window and false otherwise.

Here’s a function I wrote (amended to remove any work sensitive information):

 1from datetime import datetime, timedelta
 2
 3TIME_WINDOW_DURATION = timedelta(minutes=30)
 4
 5
 6def _check_if_user_created_in_time_window(self, account_creation):
 7    """
 8    If the user's account creation time falls within this time window,
 9    return True
10
11    Parameters
12    ----------
13        account_creation: timestamp
14            Timestamp of when user's account was created
15
16    Returns
17    -------
18    bool
19        True if tag to be applied, False otherwise
20    """
21    account_creation_datetime = self._cast_datetime_string_to_datetime_type(
22        account_creation
23    )
24    now = datetime.utcnow()
25    user_gets_tag = now - TIME_WINDOW_DURATION <= account_creation_datetime <= now
26    return user_gets_tag

Writing tests

Writing tests when you have to match against a timestamp is tricky because it could create fragile tests. In other words, a test that may or may not pass, and the pass or failure does not tell you if it is the code failing or because the timestamps do not match.

So the first tip is to use freezegun. This allows you to effectively set the date and time when the system is under test so you can make assertions against the function.

Here’s an example of this in practice:

1@freeze_time("2018-09-07 16:35:00")
2def test_user_created_in_time_window_returns_true(client):
3    example_manager = Example()
4    account_creation = "2018-09-07 16:05:01"
5
6    user_gets_tag = example_manager._check_if_user_created_in_time_window(
7        account_creation
8    )
9    assert user_gets_tag

The freeze_time decorator sets the system under test date time as 2018-09-07 16:35:00 so when we assert an account creation time of 2018-09-07 16:05:01 it falls within the time window of 30 minutes i.e. evaluates to true.

As you would expect, I wanted to make different assertions based on different frozen times and so wrote another test like the above but with a different date time passed in as the argument to the decorator. That’s all well and good as it tests the code but it goes against the DRY (don’t repeat yourself) principle.

So here’s the second trick I learned:

 1@pytest.mark.parametrize(
 2    "account_creation", ["2018-09-07 16:34:00", "2018-09-07 16:05:01"]
 3)
 4@freeze_time("2018-09-07 16:35:00")
 5def test_user_created_in_time_window_returns_true(client, account_creation):
 6    example_manager = Example()
 7
 8    user_gets_tag = example_manager._check_if_user_created_in_time_window(
 9        account_creation
10    )
11    assert user_gets_tag

Pytest - per the docs - “enables parametrization of arguments for a test function”. So how does this work?

Like freeze time, you wrap the unit test with a decorator which takes two arguments. The first is a string which is the name of the argument. This should also be passed in as an argument to the test function. The second argument is a list of the parameters. In the example above, I’ve added two different date time strings as parameters. This means when the test runs, it will run twice, using the first parameter and then the second. This keeps the code DRY whilst allowing multiple assertions. What’s also neat is when you run the tests with verbosity pytest -vv the output displays the test being run along with the parameter used. The unit test above displays:

1test_example.py::test_user_created_in_time_window_returns_true[2018-09-07 16:34:00] PASSED
2test_example.py::test_user_created_in_time_window_returns_true[2018-09-07 16:05:01] PASSED

Two nice tips to help write good unit tests and keep the code DRY.