Clean up the tests slightly