Fix TestCode

Fix TestCode