fixing test and code for python3 as well