Fix test build

Some PRs have fallen out of sync with the changes in signatures, so we
need to take a few extra parameters into account.