Fixing python unit test