Tidied up tests a bit