reset_env option to add a patch to sitecustomize.py

Author: qwcode

tests/test_pip.py

@@ -102,13 +102,22 @@ def install_setuptools(env):
 env = None
 
 
-def reset_env(environ=None, use_distribute=None, system_site_packages=False):
+def reset_env(environ=None, use_distribute=None, system_site_packages=False, sitecustomize=None):
+    """Return a test environment.
+
+    Keyword arguments:
+    environ: an environ object to use.
+    use_distribute: use distribute, not setuptools.
+    system_site_packages: create a virtualenv that simulates --system-site-packages.
+    sitecustomize: a string containing python code to add to sitecustomize.py.
+    """
+
     global env
     # FastTestPipEnv reuses env, not safe if use_distribute specified
     if use_distribute is None and not system_site_packages:
-        env = FastTestPipEnvironment(environ)
+        env = FastTestPipEnvironment(environ, sitecustomize=sitecustomize)
     else:
-        env = TestPipEnvironment(environ, use_distribute=use_distribute)
+        env = TestPipEnvironment(environ, use_distribute=use_distribute, sitecustomize=sitecustomize)
 
     if system_site_packages:
         #testing often occurs starting from a private virtualenv (e.g. with tox)
@@ -270,7 +279,7 @@ class TestPipEnvironment(TestFileEnvironment):
 
     verbose = False
 
-    def __init__(self, environ=None, use_distribute=None):
+    def __init__(self, environ=None, use_distribute=None, sitecustomize=None):
 
         self.root_path = Path(tempfile.mkdtemp('-piptest'))
 
@@ -347,6 +356,8 @@ def __init__(self, environ=None, use_distribute=None):
         # Install this version instead
         self.run('python', 'setup.py', 'install', cwd=src_folder, expect_stderr=True)
         self._use_cached_pypi_server()
+        if sitecustomize:
+            self._add_to_sitecustomize(sitecustomize)
 
     def _ignore_file(self, fn):
         if fn.endswith('__pycache__') or fn.endswith(".pyc"):
@@ -374,21 +385,30 @@ def _use_cached_pypi_server(self):
         # 'import pypi_server' ultimately imports pkg_resources (which intializes pkg_resources.working_set based on the current state of sys.path)
         # pkg_resources.get_distribution (used in pip.req) requires an accurate pkg_resources.working_set
         # therefore, 'import pypi_server' shouldn't occur in a pth file.
+
+        patch = """
+            import sys
+            sys.path.insert(0, %r)
+            import pypi_server
+            pypi_server.PyPIProxy.setup()
+            sys.path.remove(%r)""" % (str(here), str(here))
+        self._add_to_sitecustomize(patch)
+
+    def _add_to_sitecustomize(self, snippet):
+        "Adds a python code snippet to sitecustomize.py."
         sitecustomize_path = self.lib_path / 'sitecustomize.py'
-        sitecustomize = open(sitecustomize_path, 'w')
-        sitecustomize.write('import sys; ')
-        sitecustomize.write('sys.path.insert(0, %r); ' % str(here))
-        sitecustomize.write('import pypi_server; pypi_server.PyPIProxy.setup(); ')
-        sitecustomize.write('sys.path.remove(%r); ' % str(here))
+        sitecustomize = open(sitecustomize_path, 'a')
+        sitecustomize.write(textwrap.dedent('''
+                               %s
+        ''' %snippet))
         sitecustomize.close()
 
-
 fast_test_env_root = here / 'tests_cache' / 'test_ws'
 fast_test_env_backup = here / 'tests_cache' / 'test_ws_backup'
 
 
 class FastTestPipEnvironment(TestPipEnvironment):
-    def __init__(self, environ=None):
+    def __init__(self, environ=None, sitecustomize=None):
         import virtualenv
 
         self.root_path = fast_test_env_root
@@ -473,6 +493,8 @@ def __init__(self, environ=None):
             self.run('python', 'setup.py', 'install', cwd=src_folder, expect_stderr=True)
             shutil.copytree(self.root_path, self.backup_path, True)
         self._use_cached_pypi_server()
+        if sitecustomize:
+            self._add_to_sitecustomize(sitecustomize)
         assert self.root_path.exists
 
     def __del__(self):

tests/test_test.py

@@ -0,0 +1,52 @@
+"""Test the test support."""
+
+import sys
+from os.path import abspath, join, curdir, isdir, isfile
+from nose import SkipTest
+from tests.local_repos import local_checkout
+from tests.test_pip import here, reset_env, run_pip, pyversion
+
+
+patch_urlopen = """
+       def mock_urlopen():
+           pass
+       import pip
+       pip.backwardcompat.urllib2.urlopen = mock_urlopen
+    """
+
+def test_pypiproxy_patch_applied():
+    """
+    Test the PyPIProxy.setup() patch was applied, and sys.path returned to normal
+    """
+
+    env = reset_env()
+    result = env.run('python', '-c', "import pip; print(pip.backwardcompat.urllib2.urlopen.__module__)")
+    #if it were not patched, the result would be 'urllib2'
+    assert "pypi_server"== result.stdout.strip(), result.stdout
+
+    #confirm the temporary sys.path adjustment is gone
+    result = env.run('python', '-c', "import sys; print(sys.path)")
+    paths = eval(result.stdout.strip())
+    assert here not in paths, paths
+
+
+def test_add_patch_to_sitecustomize():
+    """
+    Test adding monkey patch snippet to sitecustomize.py (using TestPipEnvironment)
+    """
+
+    env = reset_env(sitecustomize=patch_urlopen, use_distribute=True)
+    result = env.run('python', '-c', "import pip; print(pip.backwardcompat.urllib2.urlopen.__module__)")
+    assert "sitecustomize"== result.stdout.strip(), result.stdout
+
+
+def test_add_patch_to_sitecustomize_fast():
+    """
+    Test adding monkey patch snippet to sitecustomize.py (using FastTestPipEnvironment)
+    """
+
+    env = reset_env(sitecustomize=patch_urlopen)
+    result = env.run('python', '-c', "import pip; print(pip.backwardcompat.urllib2.urlopen.__module__)")
+    assert "sitecustomize"== result.stdout.strip(), result.stdout
+
+