Test code cleanup