Fix tests after previous commit