Fix the tests, and fix bug found by the tests :/