Tidied up the tests a little