Fixed unit test for API refactoring