Follow API change in tests.