Modify Test class