Updated test code with comments