Yielding for testability in python
I was writing a small cli tool and wanted to try out typer. I’m a big fan of fast api, so I thought a cli library from the same author was worth a look. This blog post will cover a pattern I ended up using to help with testing.
The app and the first attempt at tests
I won’t go into too much detail about my actual problem, but I ended up with something like the following:
import typer def my_command(some_input: str): typer.echo(f"Starting up") # Do the first thing typer.echo(f"Result of first thing was X") # Do the second thing typer.echo(f"Result of second thing was X") # Do some clean up typer.echo(f"All done") if __name__ == "__main__": typer.run(my_command)
The value in the code I was writing was the output of
typer.echo. So I wanted to test it.
Initially in my tests I was monkey patching
typer.echo and asserting that it matched what
I wanted. This was okay, but I don’t like monkey patching with mocks as it often becomes
very tangled and complicated.
A solution? Yielding.
Since I already had a function representing my command invocation I thought it would be
really nice if I could just return the output. Then I could call the function and assert
that it had output what I wanted. However, given that quite significant delays could
happen in between each of my echo statements I didn’t really want to wait until the end
to get all my output. So instead
yield seemed like a good candidate. I updated my code
to look something like this:
@dataclass class Echo: message: str def my_command(some_input: str): yield Echo(f"Starting up") # Do the first thing yield Echo(f"Result of first thing was X") # Do the second thing yield Echo(f"Result of second thing was X") # Do some clean up yield Echo(f"All done")
Now my function was a generator. The code looked almost identical which was nice. I did need to add a layer on top to run this generator but this was not too complicated:
import typer def run_cli(generator_func): for output_line in generator_func(): typer.echo(output_line.message)
and now my test cases could look like this:
def test_something(): actual = list(my_command("input")) expected = [ Echo("first message"), Echo("second message"), Echo("etc.") ] assert actual == expected
I’m counting this as a success. My tests became a lot simpler which was my initial goal.
In addition, the
my_command function lost its dependency on the
typer library and became
agnostic to how the input is displayed to the user (though this wasn’t really my initial goal).