Refactor unit tests
Finish refactoring unit tests