Update tests to new version of API