[core] add sitecustomize.py test when -S is used

Author: Emanuele Palazzetti

tests/commands/ddtrace_run_sitecustomize.py

@@ -1,10 +1,19 @@
 from __future__ import print_function
 
+import sys
 from ddtrace import tracer
 from nose.tools import ok_
 
 
 if __name__ == '__main__':
+    # detect if `-S` is used
+    suppress = len(sys.argv) == 2 and sys.argv[1] is '-S'
+    if suppress:
+        ok_('sitecustomize' not in sys.modules)
+    else:
+        ok_('sitecustomize' in sys.modules)
+
+    # ensure the right `sitecustomize` will be imported
     import sitecustomize
     ok_(sitecustomize.CORRECT_IMPORT)
     print('Test success')

tests/commands/test_runner.py

@@ -5,6 +5,8 @@
 import subprocess
 import unittest
 
+from ..util import inject_sitecustomize
+
 
 class DdtraceRunTest(unittest.TestCase):
     def tearDown(self):
@@ -152,20 +154,18 @@ def test_sitecustomize_run(self):
         # [Regression test]: ensure users `sitecustomize.py` is properly loaded,
         # so that our `bootstrap/sitecustomize.py` doesn't override the one
         # defined in users' PYTHONPATH.
-        #
-        # Copy the current environment and replace the PYTHONPATH. This is
-        # required otherwise `ddtrace-run` is not found: when `env` kwarg is
-        # passed, the environment is entirely replaced
-        env = os.environ.copy()
-        sitecustomize = os.path.join(os.path.dirname(__file__), 'bootstrap')
-
-        # Add `boostrap` module so that `sitecustomize.py` is at the bottom
-        # of the PYTHONPATH
-        python_path = list(sys.path) + [sitecustomize]
-        env['PYTHONPATH'] = ':'.join(python_path)[1:]
-
+        env = inject_sitecustomize('tests/commands/bootstrap')
         out = subprocess.check_output(
             ['ddtrace-run', 'python', 'tests/commands/ddtrace_run_sitecustomize.py'],
             env=env,
         )
         assert out.startswith(b"Test success")
+
+    def test_sitecustomize_run_suppressed(self):
+        # ensure `sitecustomize.py` is not loaded if `-S` is used
+        env = inject_sitecustomize('tests/commands/bootstrap')
+        out = subprocess.check_output(
+            ['ddtrace-run', 'python', 'tests/commands/ddtrace_run_sitecustomize.py', '-S'],
+            env=env,
+        )
+        assert out.startswith(b"Test success")

tests/util.py

@@ -1,7 +1,9 @@
 import os
+import sys
 import mock
 import ddtrace
 
+from ddtrace import __file__ as root_file
 from nose.tools import ok_
 from contextlib import contextmanager
 
@@ -75,3 +77,26 @@ def set_env(**environ):
     finally:
         os.environ.clear()
         os.environ.update(old_environ)
+
+
+def inject_sitecustomize(path):
+    """Creates a new environment, injecting a ``sitecustomize.py`` module in
+    the current PYTHONPATH.
+
+    :param path: package path containing ``sitecustomize.py`` module, starting
+                 from the ddtrace root folder
+    :returns: a cloned environment that includes an altered PYTHONPATH with
+              the given `sitecustomize.py`
+    """
+    root_folder = os.path.dirname(root_file)
+    # Copy the current environment and replace the PYTHONPATH. This is
+    # required otherwise `ddtrace` scripts are not found when `env` kwarg is
+    # passed
+    env = os.environ.copy()
+    sitecustomize = os.path.join(root_folder, '..', path)
+
+    # Add `boostrap` module so that `sitecustomize.py` is at the bottom
+    # of the PYTHONPATH
+    python_path = list(sys.path) + [sitecustomize]
+    env['PYTHONPATH'] = ':'.join(python_path)[1:]
+    return env