Fix tests for py3