Refactoring unit tests with the latest changes