cleaned up code and fixed unit tests