Fixes for testing framework refactoring