Minor fixes to the tests/code