cleaned up some code and fixed broke unit test