Updated testing code