Refactor tests a bit for clarity