Refactor unit tests
Improve tests