Fix tests broken with #189

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