Refactor tests (a little)