Update test scripts to also work in Python 2