Refactor failing tests