Updated test class