Update code in the test case