Refactored and added test for search feature