code cleanup and fixes based on adding unit testing and code review