Fix unit test for python3