Add some tests

Also fix one bug found by the tests