Refactor - all tests once again pass