Adds some refactorings and tests