Some changes to test class