@contextmanager def do_run_mocked(mock): orig_run = mod.do_run mod.do_run = mock yield mod.do_run = orig_run
And then I use it in a test like this:
class TestRun(TestCase): def test_run(self): mock = MockDoRun() with do_run_mocked(mock): run(...) # This will end up calling mod.do_run # Here I can assert that run() works by inspecting what is stored in # the mock.
However, there's one thing you need to remember when using contextlib.contextmanager (and specially when using it for this purpose): if an unhandled exception occurs in the block, it is reraised inside the generator at the point where the yield occurred.
Even though it is mentioned in the docs, neither myself nor the people who reviewed my code seemed to realize its consequences -- I hope they'd call my attention to that if they did realize.
And the most important consequence (in this use case) is that when something inside your 'with' block raises an exception, the mocking won't be undone and thus the remaining tests may be affected (and fail!). Of course, once you fix the 'with' block to not raise, the mocking is undone and the remaining tests won't be affected but you'd still be left with no clue as to why they failed earlier and why they started passing again, all of a sudden. I surely wouldn't trust such test suite.
By now you might be thinking that this can be easily avoided by wrapping the yield statement with a try/finally (and do the unmocking there). Indeed it can, but I'm sure I'll keep forgetting about it so from now on I'll try to use only test fixtures (as Robert proposes) for this purpose.