Fix upgrade step tests