Update test to match new behavior