Fix testing-related code