Updating tests which broke from API refactor