modifying test class slightly