Minor modifications to test framework