fixed the unit tests
Code cleanup