Refactor - all tests passing