Minor changes to get the tests working