Code Cleanup and adding unit tests