Update the test code