Minor changes to test class