Refactored the Test class