Removed duplicate code in test class