Refactor tests so that they make slightly more sense