Yesterday I was working on a script that should run forever, or at least until the user stops it. The library behind it was already tested except for this little function:
def collect(directory): sequential = Sequential(directory) live = Live(directory) while 1: live.process() sequential.process()
this is the entry point of my library. I have an executable that just parses a directory name from command line and call this function. It’s purpose it to collect all files from a directory, filter based on some rules, and publish the file names in a queue which will be consumed by another script. It runs forever because newly created files are detected and collected too.
Anyway, what it does doesn’t really matter, the problem is: how to test this function since it’s supposed to run forever?
I don’t mind to have just unit tests for this function because I already have more integration-like tests for the classes it uses. The first solution I thought was something like this:
def collect(directory): sequential = Sequential(directory) live = Live(directory) while should_continue(): # new function to mock on tests: UGLY! live.process() sequential.process() def should_continue(): return True
this way I could mock
should_continue() in my test and make it return
False to abort the loop when I want. That works but it’s ugly! I don’t like to add dependency injections only for tests.
I asked on #python on irc and marienz gave a neat idea: raise an exception.
I could mock
sequential.process() and raise an exception, this way I know they were called as I expected and also it will also abort the loop!
import pytest import mock # this is the library under test import collectors # replace original classed with mock objects @mock.patch('collectors.Sequential') @mock.patch('collectors.Live') def test_collect_should_loop_forever_processing_both_collectors( collectors_Live, collectors_Sequential): # build mock instances. process() method will raise error # when called for the 2nd time. The code for `ErrorAfter` # is bellow seq = mock.Mock(['process']) seq.process.side_effect = ErrorAfter(2) live = mock.Mock(['process']) live.process.side_effect = ErrorAfter(2) # ensure mocked classes builds my mocked instances collectors_Sequential.return_value = seq collectors_Live.return_value = live # `ErrorAfter` will raise `CallableExhausted` with pytest.raises(CallableExhausted): collectors.collect('/tmp/files') # make sure our classed are instantiated with directory collectors_Sequential.assert_called_once_with('/tmp/files') collectors_Live.assert_called_once_with('/tmp/files')
This test uses py.test and mock. I hope the comments explains enough. The idea is simple: make
process() raise an Exception to abort the loop.
ErrorAfter class is a small helper, it builds a callable object that will raise a specific exception after
n calls. I created a custom exception here to make sure my test fails if any other exception is raised. See the code bellow.
class ErrorAfter(object): ''' Callable that will raise `CallableExhausted` exception after `limit` calls ''' def __init__(self, limit): self.limit = limit self.calls = 0 def __call__(self): self.calls += 1 if self.calls > self.limit: raise CallableExhausted class CallableExhausted(Exception): pass
Try to avoid as much as possible to create dependency injections specifically for your tests. In dynamic languages like Python it’s very easy to replace a specific component with a mock object without adding extra complexity to your code just to allow unit testing.
This was the first time I had to test a infinite loop, it’s possible and easy!