Refactored the tests to be more slick