Fix tests with the new arguments