Minor improvements to test code