Remove code duplication and add unit tests