Import pickle from compat in tests

tests/cloudpickle_file_test.py

@@ -1,7 +1,6 @@
 from __future__ import unicode_literals
 
 import os
-import pickle
 import shutil
 import sys
 import tempfile
@@ -10,6 +9,7 @@
 import pytest
 
 import cloudpickle
+from cloudpickle.compat import pickle
 
 
 class CloudPickleFileTests(unittest.TestCase):

tests/cloudpickle_test.py

@@ -9,7 +9,6 @@
 import logging
 import math
 from operator import itemgetter, attrgetter
-import pickle
 import platform
 import random
 import shutil
@@ -43,6 +42,7 @@
     tornado = None
 
 import cloudpickle
+from cloudpickle.compat import pickle
 from cloudpickle.cloudpickle import _is_dynamic
 from cloudpickle.cloudpickle import _make_empty_cell, cell_set
 from cloudpickle.cloudpickle import _extract_class_dict, _whichmodule
@@ -519,7 +519,7 @@ def test_module_locals_behavior(self):
         pickled_func_path = os.path.join(self.tmpdir, 'local_func_g.pkl')
 
         child_process_script = '''
-        import pickle
+        from cloudpickle.compat import pickle
         import gc
         with open("{pickled_func_path}", 'rb') as f:
             func = pickle.load(f)
@@ -604,7 +604,7 @@ def test_load_dynamic_module_in_grandchild_process(self):
         child_process_module_file = os.path.join(
             self.tmpdir, 'dynamic_module_from_child_process.pkl')
         child_process_script = '''
-            import pickle
+            from cloudpickle.compat import pickle
             import textwrap
 
             import cloudpickle
@@ -624,7 +624,7 @@ def test_load_dynamic_module_in_grandchild_process(self):
 
         # The script ran by the process created by the child process
         child_of_child_process_script = """ '''
-                import pickle
+                from cloudpickle.compat import pickle
                 with open('{child_process_module_file}','rb') as fid:
                     mod = pickle.load(fid)
                 ''' """
@@ -679,7 +679,7 @@ def my_small_function(x, y):
         assert b'math' not in b
 
     def test_is_dynamic_module(self):
-        import pickle  # decouple this test from global imports
+        from cloudpickle.compat import pickle
         import os.path
         import distutils
         import distutils.ccompiler
@@ -969,7 +969,8 @@ def example():
 
         # choose "subprocess" rather than "multiprocessing" because the latter
         # library uses fork to preserve the parent environment.
-        command = ("import pickle, base64; "
+        command = ("import base64; "
+                   "from cloudpickle.compat import pickle; "
                    "pickle.loads(base64.b32decode('" +
                    base64.b32encode(s).decode('ascii') +
                    "'))()")
@@ -990,7 +991,8 @@ def example():
 
         s = cloudpickle.dumps(example, protocol=self.protocol)
 
-        command = ("import pickle, base64; "
+        command = ("import base64; "
+                   "from cloudpickle.compat import pickle; "
                    "pickle.loads(base64.b32decode('" +
                    base64.b32encode(s).decode('ascii') +
                    "'))()")

tests/testutils.py

@@ -4,14 +4,15 @@
 import tempfile
 import base64
 from subprocess import Popen, check_output, PIPE, STDOUT, CalledProcessError
-from pickle import loads
+from cloudpickle.compat import pickle
 from contextlib import contextmanager
 from concurrent.futures import ProcessPoolExecutor
 
 import psutil
 from cloudpickle import dumps
 from subprocess import TimeoutExpired
 
+loads = pickle.loads
 TIMEOUT = 60
 TEST_GLOBALS = "a test value"