Updated test code