fix testing code