Fix Test code