Fix tests related to last big refactor.