Refactoring and writing tests