Debugging Your Unit Test Suite in Python

Fix errors in your unit tests and mock out components with ease - how to use the Python debugger pdb in pytest


Have you ever found it harder to write the unit tests than the actual business logic? Python has a reputation as being a simple language to work in, but that doesn’t always extend to the unit tests; some things are just really awkward to write tests for. Perhaps the function you’re testing requires you to mock out a bunch of things, or perhaps it outputs a complex data structure like a DataFrame and takes a ton of assertions to properly specify the desired behaviour.

In any case, these types of tests aren’t much fun to write whether you’re a software engineer or a data scientist. Writing unit tests can be tedious work at the best of times, but with added difficulties like mocks it can take forever to get them done. In the worst case, you might find yourself doing programming by trial-and-error--you think your code is right but your tests are failing, so you just keep tweaking the test code and re-running until it passes.

Needless to say, trial-and-error is a terrible methodology for writing tests. It takes a long time and might bake some bugs into the test suite rather than preventing them. Luckily, there is an extremely versatile tool that makes this easier: the Python debugger pdb. And if you’re using pytest to run your tests, it has a handy integration with pdb.

In this article I am going to demonstrate how to use pdb with pytest. First I’ll give a short introduction to debuggers and how they compare with other tools like Jupyter. Then I’ll move to a working example of a unit test that mocks a database connection. The test code has a confusing error, but we can drop into the debugger to pinpoint exactly what’s wrong and resolve the issue--a hard problem made easy.

So with that, let’s get started!

A high level case for using debuggers

Pdb is the basic debugger that comes with Python. It lets you pause a running program and inspect the values of variables, print things out, and even make live changes. If you’ve never used pdb, or even seen the need for a debugger in an interpreted language like Python, its killer feature is this -- when an exception is encountered, you can drop into an interactive environment instead of crashing the program, also known as post-mortem debugging. This lets you see exactly what's broken and what's needed, turning a frustrating trial-and-error experience into a fast and easy exercise.

If you’ve done work in Jupyter or a similar environment, you may already be used to printing stuff out and poking around after an error, but pdb takes it one step further. While Jupyter drops you back at the top level of your code, pdb pauses at the exact line where the exception was raised--a priceless feature when the error happens deep in a function call.

Let’s say you call a function, which calls another function, which calls another function, which hits an error with a confusing message, perhaps something deep within Python’s core or a library like Pandas. The actual bug is probably in one of those intermediate steps, not at the top (your initial line of code) nor the bottom (library code). That’s why pdb also has the handy commands up and down, letting you navigate the call stack and check out all the local variables in each level of the error traceback.

For a more detailed introduction to pdb, I recommend Nathan Jennings’s pdb tutorial at Real Python. And for the Jupyter aficionados out there, try running %debug in the next cell after an exception.

Using pytest --pdb

If you run your tests with pytest --pdb, it will automatically drop into a debugger on every test that fails or has an error of some kind. If you love using the REPL, IPython or Jupyter in your normal work, then you know how handy it is to be in an interactive environment when trying to fix errors! Pytest has many helpful features for writing unit tests, but this is among the most versatile.

If you want to start debugging at a specific point, you can also add a breakpoint() line (available in Python 3.7+) right before the point of interest.

Once you’re there, you can use all the usual debugger commands - break, step, next, return, up, down, continue, plus whatever arbitrary Python code you want to execute.

Ctrl+D or exit will let you quit the debugger.

Example -- debugging a unit test with a mocked out database

One place where I find using pdb really useful is when I’m writing tests with mocked components. For instance, I might have some function that queries a database in order to do some calculations. I don’t actually want to connect to the database in my unit tests since it would be slow and require a username and password. Instead, I can mock out the database connection and simulate the results of the queries.

Writing the business logic and test case

Let’s see what that might look like in code--feel free to copy and paste these snippets and run them on your computer if you want to follow along.

First, let’s sketch out our actual function. We’re using the Capital One open source project locopy to connect to our database, which provides a handy ContextManager that we can use with the with keyword to get a connection that cleans up after itself. Just in case you don’t have locopy installed, we can define a stub class for Database which will allow the tests to run for now.

    # my_pkg.py
try:
   from locopy import Database
except (ImportError, ModuleNotFoundError):
   class Database: pass

connection_params = {}

def refresh_views(purge_expired: bool = False):
   with Database(**connection_params) as db:
       db.execute("query1")
       db.execute("query2")

       if purge_expired:
           db.execute("query3")
  

So our function connects to the database and runs some queries (query1 and query2). If we set the purge_expired flag, it also runs query3.

Now let’s write a test case to ensure that the purge_expired option works as expected.

    # test_my_pkg.py
from unittest.mock import patch
import pytest
from my_pkg import refresh_views

@patch("my_pkg.Database")
def test_refresh_views(mock_Database):
   mock_db = mock_Database()

   # make sure query3 isn't run
   refresh_views(purge_expired=False)
   with pytest.raises(AssertionError):
       mock_db.execute.assert_called_with("query3")

   # make sure query3 _is_ run
   refresh_views(purge_expired=True)
   mock_db.execute.assert_called_with("query3")
  

Using @patch from unittest.mock, we can replace the Database class with a MagicMock instance while the test runs. If you’re unfamiliar with it, MagicMock provides a dummy object that will pretend to have whatever methods or attributes you need (though they don’t actually do anything by default). In your tests, you can inspect the mock to see which of its methods were called or supply test data that it will return on your behalf.

In this case, we are checking to see that query3 doesn’t get executed by the database when we set purge_expired=False, and that it does when we set purge_expired=True.

Mocks are really handy, but MagicMock is so unlike normal Python objects that programming with it can be quite confusing. This is especially true since we normally only use mocks in a unit testing environment, so most of us don’t have a lot of hands-on experience with them.

Pay special attention to the way we define mock_db, which is supposed to match the variable db inside refresh_views. @patch lets us replace the Database class itself with mock_Database but we’re not using the class directly; we’re using an instance of it. That means we need to “instantiate” our mock, hence mock_db = mock_Database(). There isn’t any special significance to this except that every transformation we make in our original code, we should also apply to the mock.

Now let’s run our test and… we get an error.

    @patch("my_pkg.Database")
   def test_refresh_views(mock_Database):
       mock_db = mock_Database()

       # make sure query3 isn't run
       refresh_views(purge_expired=False)
       with pytest.raises(AssertionError):
           mock_db.execute.assert_called_with("query3")

       # make sure query3 _is_ run
       refresh_views(purge_expired=True)
>      mock_db.execute.assert_called_with("query3")

E           AssertionError: expected call not found.
E           Expected: execute('query3')
E           Actual: not called.
  

That’s strange.

Debugging the test case

So, what’s our next step? There are a lot of things that can go wrong with mocks - it could be that query3 never got executed (i.e., the test legitimately failed), that the patch didn’t work properly, or perhaps we didn’t perfectly replicate all the steps between Database and db on our mock.

My first step in this situation is to rerun the test with the --pdb option and start looking around.

    ~/miniconda3/lib/python3.9/unittest/mock.py:898: AssertionError
>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>> entering PDB >>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>

>>>>>>>>>>>>>>>>>> PDB post_mortem (IO-capturing turned off) >>>>>>>>>>>>>>>>>>>
> ~/miniconda3/lib/python3.9/unittest/mock.py(898)assert_called_with()
-> raise AssertionError(error_message)
(Pdb)
  

The traceback starts inside of a library function, so I execute up once to get back to my test code.

    (Pdb) up
> ~/dev/test_my_pkg.py(17)test_refresh_views()
-> mock_db.execute.assert_called_with("query3")
  

Okay, now we can inspect mock_db a bit.

    (Pdb) mock_db


(Pdb) mock_db.execute


(Pdb)
  

We can see mock_db is supposedly a Database instance. Curiously the execute method seems as though it was never called since it has an empty call list. Exiting the debugger for a moment, let’s try again from inside the actual implementation code. I’ll set a breakpoint right before query3 is supposed to be executed.

    def refresh_views(purge_expired: bool = False):
   with Database(**connection_params) as db:
       db.execute("query1")
       db.execute("query2")

       if purge_expired:
>          breakpoint()
           db.execute("query3")
  

Now when we run the tests, we drop into the debugger on this line. The first step is to see what db looks like from the inside - is it even a mock at all, and if so, what methods have been called on it?

    > ~/dev/my_pkg.py(16)refresh_views()
-> db.execute("query3")

(Pdb) db

  

Right away, it jumps out that db is a mock, but it’s not just Database() but Database().__enter__(). Where did that come from?

Well, if we look at where it’s defined, we see that it comes from a with statement - that means it’s actually a ContextManager.

    with Database(**connection_params) as db:
  

Personally, I use with statements all the time in Python but usually don’t create my own custom ContextManagers. So while I initially suspected the with statement I couldn’t remember the exact syntax to replicate with the mock.

Applying the fix

Now that we know the problem, we can fix our test code. We could manually replicate the chain of method calls on Database...

    @patch("my_pkg.Database")
def test_refresh_views(mock_Database):
>  mock_db = mock_Database().__enter__()

   # make sure query3 isn't run
   refresh_views(purge_expired=False)
   ...
  

Alternatively, we could define it with with, the exact same way as in the business logic. Very intuitive.

    @patch("my_pkg.Database")
def test_refresh_views(mock_Database):
>  with mock_Database() as mock_db:
       # make sure query3 isn't run
       refresh_views(purge_expired=False)
  

After that small edit (and removal of breakpoint()), the test passes!

    ========================= 1 passed in 0.64s =========================
  

If you had asked me a few years ago I might not have even known what to Google for this problem. But with the debugger I was able to solve the issue right away with no fuss. Being able to self-serve is a wonderful thing!

Wrap up

It’s much easier to fix bugs when you can inspect every variable for yourself, and that applies just as much to test code as it does to everything else. But even if you’re used to poking around in Jupyter, it’s not obvious how to apply those same techniques in unit testing. Not to mention, MagicMock behaves so differently from normal Python objects that it’s quite easy to get mixed up. That’s what pytest --pdb and breakpoint() are for.

So say goodbye to trial-and-error programming and endless Googling, at least when writing unit tests. With this trick, your most frustrating test bugs should be a piece of cake. 🎂


Robin Neufeld, Data Scientist

Robin Neufeld is a natural scientist turned data scientist. She loves all things Python, studying (human) languages, and winning fights with the computer.


DISCLOSURE STATEMENT: © 2021 Capital One. Opinions are those of the individual author. Unless noted otherwise in this post, Capital One is not affiliated with, nor endorsed by, any of the companies mentioned. All trademarks and other intellectual property used or displayed are property of their respective owners.

Related Content