Adding junit libs
Added junit test
Cleaned up some code