Chaining ContextManager Values When Mocking Tests
Update — Mar 1st, 2021
Hello! If you’ve seen my articles pop up in your searches before, I really appreciate the time people have taken to read my articles. At this point, there’s been over 75k views across the 25 or so articles I’ve written here on Medium. That’s amazing to me since writing is a hobby I started doing just to try to get a bit better at writing and communicating.
I wanted to let everyone who lands on these articles that I have migrated everything over to a self hosted blog (for a bunch of reasons that I won’t get into here). So please take a look at https://holdenrehg.com/blog for all of the articles or https://holdenrehg.com/blog/2020-01-29_chaining-context-managers-for-mocks to specifically see this article.
Thanks again! — Holden
Recently, I’ve used unittest.mock.patch to mock connections to external services while developing a Django app. While writing the tests, I found myself, repeatedly writing code that looks like:
It’s not a huge deal at first, but as you add more use cases things really bloat up. I wanted utilities with a clear interface for the team to prevent mistakes, to make the tests more readable, and to save a few characters. Generally, the solution to, either with variables or functions, dynamically chain together context managers wasn’t perfectly clear (at least it wasn’t to me).
If you aren’t familiar with context managers, go read through the data model docs and the language reference docs on the with statement.
Variables
One potential solution we have is to pull out some common variables.
Looking back at our example, we may want to split out the authenticators so that we can easily handle both authenticated and unauthenticated scenarios.
I’m going to do that by extending the MockAuthenticator to take in a parameter called authenticated .
Then we can pull out some variables for each scenario to re-use throughout the entire test suite. This actually isn’t a bad option.
But I didn’t really like the fact that I had to redefine test_connection on every test, no matter what type of connection I was using. I attempted to abstract those up into a single context manager, but ran into issues.
The with statement can handle multiple “arguments” as you see in the example above, but you can’t pass an iterable or unpack an iterable the way you can with function calls. That would make things simpler because we would only be dealing with a couple of mock.patch variables instantiated within a reusable authed function.
Yield functions
The best option in my mind is to pull out two functions that make it clear what is going on in each test, while mocking multiple objects behind the scenes.
We can do this by defining contextmanager functions, running each mock as needed and then yield ing to the calling function. The contextmanager function is a way to define a context manager without needing to have explicit __enter__ and __exit__ methods like you would in a class definition.
Now we redefine our tests. Check them out. I personally think they are much easier to understand at a glance, with a simple interface for auth mocks. Plus it will save us a few characters as an added bonus.
A note about exit stacks
While it doesn’t make sense for the exact scenario we’re dealing with above, we could have also potentially used an ExitStack .
I had no clue what an ExitStack was until I started hunting around for a way to chain together context managers dynamically. In the python docs, it’s defined as “[a way] to make it easy to programmatically combine other context managers and cleanup functions, especially those that are optional or otherwise driven by input data.”
Sounds great. The general idea behind them is that I can enter into the ExitStack and have more control over the context managers that I push onto the stack. When the ExitStack closes then it iteratively pops every registered context manager off the stack and closes them in reverse order of their registration (first in last out).
Check out what that looks like:
This is essentially the same as running nested with statements except you’ll see how we need it to abstract out an authed and unauthed context manager.
The primary benefit is that while you cannot, for example, iterate over a list of files and pass a list of open results into a with statement, we can do that with the ExitStack .
We can take a look at another common example when dealing with mocks. Assume that you need to test a connection to an external API or database and you want to create a single test that interacts with a number of outside mocks.
You might start with something like this for each service adapter:
Instead, we can create a context manager that iterates over all possible adapters, creates a new context for them via a db.connect function, and then yields those back to the test. Once that test has finished running, then the ExitStack handles all of the cleanup for those connections.
Context managers can be an extremely useful abstraction in your python code if you don’t go overboard with them. Give it a shot in your code or tests to see where that might make sense.
Thanks For Reading
Sign up for my new ERP project called Buster and let me know your thoughts.
🐦 Follow me on Twitter for random thoughts and articles about business software and software development.
📓 Follow me directly here at Medium to stay up to date with my articles.
📷 Follow me on Instagram if you want to see pictures of (mostly) food and bread that I cook.
Holden
