remove unused code
Adding unit tests