Test code clean up