Updated test code to do testing