Refactored some code for better organization
Fixed some unit tests