Minor tweaks on test script