- Refactor unit tests