Minor updates to test framework