[testing] Ensure consistent use of override_config and override_env (#815)

* [testing] Ensure consistent use of override_config/override_env

* fix linting issues

* fix usage of override_env

* fix override_env logic

Author: brettlangdon

tests/base/__init__.py

@@ -1,4 +1,5 @@
 import contextlib
+import os
 import unittest
 
 import ddtrace
@@ -23,6 +24,25 @@ def test_case(self):
                     pass
     """
 
+    @contextlib.contextmanager
+    def override_env(self, env):
+        """
+        Temporarily override ``os.environ`` with provided values
+        >>> with self.override_env(dict(DATADOG_TRACE_DEBUG=True)):
+            # Your test
+        """
+        # Copy the full original environment
+        original = dict(os.environ)
+
+        # Update based on the passed in arguments
+        os.environ.update(env)
+        try:
+            yield
+        finally:
+            # Full clear the environment out and reset back to the original
+            os.environ.clear()
+            os.environ.update(original)
+
     @contextlib.contextmanager
     def override_config(self, integration, values):
         """

tests/contrib/celery/base.py

@@ -1,6 +1,6 @@
 from celery import Celery
 
-from ddtrace import Pin, config
+from ddtrace import Pin
 from ddtrace.compat import PY2
 from ddtrace.contrib.celery import patch, unpatch
 
@@ -21,8 +21,6 @@ class CeleryBaseTestCase(BaseTracerTestCase):
     def setUp(self):
         super(CeleryBaseTestCase, self).setUp()
 
-        # keep track of original config
-        self._config = dict(config.celery)
         # instrument Celery and create an app with Broker and Result backends
         patch()
         self.pin = Pin(service='celery-unittest', tracer=self.tracer)
@@ -34,9 +32,6 @@ def tearDown(self):
         # remove instrumentation from Celery
         unpatch()
         self.app = None
-        # restore the global configuration
-        config.celery.update(self._config)
-        self._config = None
 
         super(CeleryBaseTestCase, self).tearDown()
 

tests/contrib/celery/test_integration.py

@@ -3,7 +3,6 @@
 
 from nose.tools import eq_, ok_
 
-from ddtrace import config
 from ddtrace.contrib.celery import patch, unpatch
 
 from .base import CeleryBaseTestCase
@@ -306,41 +305,39 @@ def add(x, y):
         eq_(span.get_tag('celery.state'), 'SUCCESS')
 
     def test_worker_service_name(self):
-        # Ensure worker service name can be changed via
-        # configuration object
-        config.celery['worker_service_name'] = 'worker-notify'
-
         @self.app.task
         def fn_task():
             return 42
 
-        t = fn_task.apply()
-        ok_(t.successful())
-        eq_(42, t.result)
+        # Ensure worker service name can be changed via
+        # configuration object
+        with self.override_config('celery', dict(worker_service_name='worker-notify')):
+            t = fn_task.apply()
+            self.assertTrue(t.successful())
+            self.assertEqual(42, t.result)
 
-        traces = self.tracer.writer.pop_traces()
-        eq_(1, len(traces))
-        eq_(1, len(traces[0]))
-        span = traces[0][0]
-        eq_(span.service, 'worker-notify')
+            traces = self.tracer.writer.pop_traces()
+            self.assertEqual(1, len(traces))
+            self.assertEqual(1, len(traces[0]))
+            span = traces[0][0]
+            self.assertEqual(span.service, 'worker-notify')
 
     def test_producer_service_name(self):
-        # Ensure producer service name can be changed via
-        # configuration object
-        config.celery['producer_service_name'] = 'task-queue'
-
         @self.app.task
         def fn_task():
             return 42
 
-        t = fn_task.delay()
-        eq_('PENDING', t.status)
-
-        traces = self.tracer.writer.pop_traces()
-        eq_(1, len(traces))
-        eq_(1, len(traces[0]))
-        span = traces[0][0]
-        eq_(span.service, 'task-queue')
+        # Ensure producer service name can be changed via
+        # configuration object
+        with self.override_config('celery', dict(producer_service_name='task-queue')):
+            t = fn_task.delay()
+            self.assertEqual('PENDING', t.status)
+
+            traces = self.tracer.writer.pop_traces()
+            self.assertEqual(1, len(traces))
+            self.assertEqual(1, len(traces[0]))
+            span = traces[0][0]
+            self.assertEqual(span.service, 'task-queue')
 
     def test_fn_task_apply_async_ot(self):
         """OpenTracing version of test_fn_task_apply_async."""

tests/contrib/django/test_instrumentation.py

@@ -6,7 +6,6 @@
 
 # testing
 from .utils import DjangoTraceTestCase
-from ...util import set_env
 
 
 class DjangoInstrumentationTest(DjangoTraceTestCase):
@@ -23,17 +22,17 @@ def test_tracer_flags(self):
     def test_environment_vars(self):
         # Django defaults can be overridden by env vars, ensuring that
         # environment strings are properly converted
-        with set_env(
+        with self.override_env(dict(
             DATADOG_TRACE_AGENT_HOSTNAME='agent.consul.local',
             DATADOG_TRACE_AGENT_PORT='58126'
-        ):
+        )):
             settings = DatadogSettings()
             eq_(settings.AGENT_HOSTNAME, 'agent.consul.local')
             eq_(settings.AGENT_PORT, 58126)
 
     def test_environment_var_wrong_port(self):
         # ensures that a wrong Agent Port doesn't crash the system
         # and defaults to 8126
-        with set_env(DATADOG_TRACE_AGENT_PORT='something'):
+        with self.override_env(dict(DATADOG_TRACE_AGENT_PORT='something')):
             settings = DatadogSettings()
             eq_(settings.AGENT_PORT, 8126)

tests/contrib/falcon/test_suite.py

@@ -5,7 +5,6 @@
 from ddtrace.ext import errors as errx, http as httpx
 
 from tests.opentracer.utils import init_tracer
-from ...util import override_config
 
 
 class FalconTestCase(object):
@@ -193,7 +192,7 @@ def on_falcon_request(span, request, response):
         eq_(span.get_tag('my.custom'), 'tag')
 
     def test_http_header_tracing(self):
-        with override_config('falcon', {}):
+        with self.override_config('falcon', {}):
             config.falcon.http.trace_headers(['my-header', 'my-response-header'])
             self.simulate_get('/200', headers={'my-header': 'my_value'})
             traces = self.tracer.writer.pop_traces()

tests/contrib/flask/test_request.py

@@ -5,7 +5,6 @@
 from ddtrace.propagation.http import HTTP_HEADER_TRACE_ID, HTTP_HEADER_PARENT_ID
 from flask import abort
 
-from ...util import override_config
 from . import BaseFlaskTestCase
 
 
@@ -117,7 +116,7 @@ def index():
             return 'Hello Flask', 200
 
         # Enable distributed tracing
-        with override_config('flask', dict(distributed_tracing_enabled=True)):
+        with self.override_config('flask', dict(distributed_tracing_enabled=True)):
             res = self.client.get('/', headers={
                 HTTP_HEADER_PARENT_ID: '12345',
                 HTTP_HEADER_TRACE_ID: '678910',
@@ -131,7 +130,7 @@ def index():
         self.assertEqual(span.parent_id, 12345)
 
         # With distributed tracing disabled
-        with override_config('flask', dict(distributed_tracing_enabled=False)):
+        with self.override_config('flask', dict(distributed_tracing_enabled=False)):
             res = self.client.get('/', headers={
                 HTTP_HEADER_PARENT_ID: '12345',
                 HTTP_HEADER_TRACE_ID: '678910',

tests/contrib/httplib/test_httplib.py

@@ -1,7 +1,6 @@
 # Standard library
 import contextlib
 import sys
-import unittest
 
 # Third party
 import wrapt
@@ -15,8 +14,8 @@
 
 from tests.opentracer.utils import init_tracer
 
-from ...test_tracer import get_dummy_tracer
-from ...util import assert_dict_issuperset, override_global_tracer, override_config
+from ...base import BaseTracerTestCase
+from ...util import assert_dict_issuperset, override_global_tracer
 
 if PY2:
     from urllib2 import urlopen, build_opener, Request
@@ -39,16 +38,19 @@ def to_str(self, value):
         return value.decode('utf-8')
 
     def setUp(self):
+        super(HTTPLibBaseMixin, self).setUp()
+
         patch()
-        self.tracer = get_dummy_tracer()
         Pin.override(httplib, tracer=self.tracer)
 
     def tearDown(self):
         unpatch()
 
+        super(HTTPLibBaseMixin, self).tearDown()
+
 
 # Main test cases for httplib/http.client and urllib2/urllib.request
-class HTTPLibTestCase(HTTPLibBaseMixin, unittest.TestCase):
+class HTTPLibTestCase(HTTPLibBaseMixin, BaseTracerTestCase):
     SPAN_NAME = 'httplib.request' if PY2 else 'http.client.request'
 
     def to_str(self, value):
@@ -65,13 +67,6 @@ def get_https_connection(self, *args, **kwargs):
         Pin.override(conn, tracer=self.tracer)
         return conn
 
-    def setUp(self):
-        patch()
-        self.tracer = get_dummy_tracer()
-
-    def tearDown(self):
-        unpatch()
-
     def test_patch(self):
         """
         When patching httplib
@@ -354,7 +349,7 @@ def test_httplib_request_and_response_headers(self):
             self.assertEqual(s.get_tag('http.response.headers.access_control_allow_origin'), None)
 
         # Enabled when configured
-        with override_config('hhtplib', {}):
+        with self.override_config('hhtplib', {}):
             integration_config = config.httplib  # type: IntegrationConfig
             integration_config.http.trace_headers(['my-header', 'access-control-allow-origin'])
             conn = self.get_http_connection(SOCKET)
@@ -502,7 +497,7 @@ def test_httplib_request_get_request_ot(self):
 if PY2:
     import urllib
 
-    class HTTPLibPython2Test(HTTPLibBaseMixin, unittest.TestCase):
+    class HTTPLibPython2Test(HTTPLibBaseMixin, BaseTracerTestCase):
         def test_urllib_request(self):
             """
             When making a request via urllib.urlopen

tests/contrib/molten/test_molten.py

@@ -11,7 +11,6 @@
 from ddtrace.contrib.molten.patch import MOLTEN_VERSION
 
 from ...base import BaseTracerTestCase
-from ...util import override_config
 
 
 # NOTE: Type annotations required by molten otherwise parameters cannot be coerced
@@ -158,7 +157,7 @@ def test_resources(self):
 
     def test_distributed_tracing(self):
         """ Tests whether span IDs are propogated when distributed tracing is on """
-        with override_config('molten', dict(distributed_tracing=True)):
+        with self.override_config('molten', dict(distributed_tracing=True)):
             response = molten_client(headers={
                 HTTP_HEADER_TRACE_ID: '100',
                 HTTP_HEADER_PARENT_ID: '42',
@@ -173,7 +172,7 @@ def test_distributed_tracing(self):
         self.assertEqual(span.parent_id, 42)
 
         # Now without tracing on
-        with override_config('molten', dict(distributed_tracing=False)):
+        with self.override_config('molten', dict(distributed_tracing=False)):
             response = molten_client(headers={
                 HTTP_HEADER_TRACE_ID: '100',
                 HTTP_HEADER_PARENT_ID: '42',

tests/contrib/requests/test_requests.py

@@ -1,5 +1,3 @@
-import unittest
-
 import requests
 from requests import Session
 from requests.exceptions import MissingSchema
@@ -11,30 +9,33 @@
 
 from tests.opentracer.utils import init_tracer
 
-from ...test_tracer import get_dummy_tracer
-from ...util import override_global_tracer, override_config
+from ...base import BaseTracerTestCase
+from ...util import override_global_tracer
 
 # socket name comes from https://english.stackexchange.com/a/44048
 SOCKET = 'httpbin.org'
 URL_200 = 'http://{}/status/200'.format(SOCKET)
 URL_500 = 'http://{}/status/500'.format(SOCKET)
 
 
-class BaseRequestTestCase(unittest.TestCase):
+class BaseRequestTestCase(object):
     """Create a traced Session, patching during the setUp and
     unpatching after the tearDown
     """
     def setUp(self):
+        super(BaseRequestTestCase, self).setUp()
+
         patch()
-        self.tracer = get_dummy_tracer()
         self.session = Session()
         setattr(self.session, 'datadog_tracer', self.tracer)
 
     def tearDown(self):
         unpatch()
 
+        super(BaseRequestTestCase, self).tearDown()
+
 
-class TestRequests(BaseRequestTestCase):
+class TestRequests(BaseRequestTestCase, BaseTracerTestCase):
     def test_resource_path(self):
         out = self.session.get(URL_200)
         eq_(out.status_code, 200)
@@ -375,7 +376,7 @@ def test_request_and_response_headers(self):
         eq_(s.get_tag('http.response.headers.access-control-allow-origin'), None)
 
         # Enabled when explicitly configured
-        with override_config('requests', {}):
+        with self.override_config('requests', {}):
             config.requests.http.trace_headers(['my-header', 'access-control-allow-origin'])
             self.session.get(URL_200, headers={'my-header': 'my_value'})
             spans = self.tracer.writer.pop()

tests/contrib/requests/test_requests_distributed.py

@@ -3,10 +3,11 @@
 
 from ddtrace import config
 
+from ...base import BaseTracerTestCase
 from .test_requests import BaseRequestTestCase
 
 
-class TestRequestsDistributed(BaseRequestTestCase):
+class TestRequestsDistributed(BaseRequestTestCase, BaseTracerTestCase):
     def headers_here(self, tracer, request, root_span):
         # Use an additional matcher to query the request headers.
         # This is because the parent_id can only been known within such a callback,