Update code for test