Update unit tests for the latest API