Fixing test in python 3