Clean up traceback
Correct tests