refactored and simplified code and tests