Fix test to go with previous commit