Fix test to work with all versions of Python