Changed test to reflect code