Minor bug fixes to test bench