Refactor unit tests for consistency