Fixed unit test
Refactored some code