Refactor and fix test errors