Fix test for Python2