code cleanup and starting to unit test