Removed duplicate code
Adding tests