TST: Fix broken tests that got merged in #139

Sorry team!

![](http://i.imgur.com/YIvN7NP.gif)