Minor style fixes to test.py