Refactor tests some more