The test code generally should be refactored for robustness and clarity. cc @devin-petersohn @simon-mo @pschafhalter