Slight tweak to test code