Refactoring
First round of unit tests