Update mocks to reflect recent API changes