How to Limit Test Time in Django’s Test Framework

2021-01-25

I recently optimized a client project’s test suite, and in the process found a test whose runtime had crept up ever since it had been written. The problematic test exercised an import process from a fixed past date until the current day. The test’s runtime therefore grew every day, until it reached over a minute.

The solution for that test was to add a date mock using time-machine, so the test created a fixed amount of data. But we also wanted to prevent such a problem occurring again.

We came up with the idea of implementing a test time limit that would automatically fail any long running tests. This way, if any appeared again, they would be detected early.

The project uses Django’s test framework, which is based on unittest. Based on my spelunking of the unittest’s internals, I figured the best place to add such a limit would be in the TestCase class. The project already had its own customized subclasses of Django’s TestCase classes, so I could add the logic there.

Cutting out other details, here is the implementation I came up with:

import time

from django import test


class SlowTestException(Exception):
    pass


class ProjectTestCaseMixin:
    def _callTestMethod(self, method):
        start = time.time()

        result = super()._callTestMethod(method)

        limit_seconds = 10
        time_taken = time.time() - start
        if time_taken > limit_seconds:
            raise SlowTestException(
                f"This test took {time_taken:.2f}s, more than the limit of {limit_seconds}s."
            )

        return result


class SimpleTestCase(ProjectTestCaseMixin, test.SimpleTestCase):
    pass


class TransactionTestCase(ProjectTestCaseMixin, test.TransactionTestCase):
    pass


class TestCase(ProjectTestCaseMixin, test.TestCase):
    pass

The mixin wraps the unittest internal _callTestMethod(), which is called for each test, with the time check. Regardless of whether a test passes or fails, if it takes longer than the limit, a SlowTestException is raised, which fails the test with an error.

I checked what this looks like with a dummy slow test that runs time.sleep(11):

$ ./manage.py test example.tests.test_sleepy.SleepyTest
System check identified no issues (0 silenced).
E
======================================================================
ERROR: test_that_i_can_take_a_nap (example.tests.test_sleepy.SleepyTest)
----------------------------------------------------------------------
Traceback (most recent call last):
  File "/.../example/test.py", line 47, in _callTestMethod
    raise SlowTestException(
example.test.SlowTestException: This test took 11.00s, more than the limit of 10s.

----------------------------------------------------------------------
Ran 1 test in 11.001s

FAILED (errors=1)

Nice.

Parallel Problems

One thing this started to reveal was that when running tests in parallel mode, there would often be a handful of slow tests. But when run individually or in non-parallel mode, those tests didn’t come near the time limit.

This seemed to be cause by resource starvation on the CPU or inside PostgreSQL. The solution was to reduce the default number of parallel test processes from 8 to 4.

Tests stopped hitting the time limit and the time for the whole test run was reduced a bit, which was nice to see.

Timing Out

Setting a limit like this is simple and will prevent slow tests being added to the project. But it still waits for the whole, slow test to complete before failing it. This would be problematic for tests that could sometimes take very long to complete.

Whilst it’s best to avoid creating such tests in the first place, if you have them you might want a fuller timeout solution that stops test as soon as they hit the limit. This is possible using a thread or, on Unix, requesting an interrupt from the SIGALARM signal. The pytest-timeout implements both these methods and some code can probably be copied from there to unittest-based projects.

Fin

May your test suite ever expand its coverage and reduce its duration,

—Adam

Want better tests? Check out my book Speed Up Your Django Tests which teaches you to write faster, more accurate tests.

One summary email a week, no spam, I pinky promise.

Related posts:

Tags: django, python