Fix tests broken with #189

@abourget I hope you don't mind me fixing tests for your pull request but I need to progress on this issue :)