A minor fix to the test bench