added new test and refactored the rest