Code cleanup in test suite