Fixed failing unit tests

I hope this fixes the failing unit tests. It works locally, let's see what travis thinks of it.