requests: Improvements for requests integration (#573)

mauriciovasquezbernal · web-flow · commit 44b592cf2dff · 2020-04-14T20:20:54.000-07:00
Adding a TestBase class which wraps a tracer provider that is configured with a memory span
exporter.  This class inherits from unitest.TestCase, hence other test classes
can inherit from it to get access to the underlying memory span exporter and
tracer provider.

Adding a mock propagator that could be used for testing propagation in different packages.
It was implemented in the opentracing-shim and this commit moves it to a generic
place.

Adding disable_session(), which can be used to disable the instrumentation on a single
requests' session object.
diff --git a/dev-requirements.txt b/dev-requirements.txt
@@ -9,3 +9,4 @@ sphinx-autodoc-typehints~=1.10.2
 pytest!=5.2.3
 pytest-cov>=2.8
 readme-renderer~=24.0
+httpretty~=1.0
diff --git a/ext/opentelemetry-ext-http-requests/README.rst b/ext/opentelemetry-ext-http-requests/README.rst
@@ -7,7 +7,7 @@ OpenTelemetry requests Integration
    :target: https://pypi.org/project/opentelemetry-ext-http-requests/
 
 This library allows tracing HTTP requests made by the
-`requests <https://requests.kennethreitz.org/en/master/>`_ library.
+`requests <https://requests.readthedocs.io/en/master/>`_ library.
 
 Installation
 ------------
diff --git a/ext/opentelemetry-ext-http-requests/setup.cfg b/ext/opentelemetry-ext-http-requests/setup.cfg
@@ -43,5 +43,10 @@ install_requires =
     opentelemetry-api == 0.7.dev0
     requests ~= 2.0
 
+[options.extras_require]
+test =
+    opentelemetry-test == 0.7.dev0
+    httpretty ~= 1.0
+
 [options.packages.find]
 where = src
diff --git a/ext/opentelemetry-ext-http-requests/src/opentelemetry/ext/http_requests/__init__.py b/ext/opentelemetry-ext-http-requests/src/opentelemetry/ext/http_requests/__init__.py
@@ -42,6 +42,7 @@
 """
 
 import functools
+import types
 from urllib.parse import urlparse
 
 from requests.sessions import Session
@@ -96,9 +97,6 @@ def instrumented_request(self, method, url, *args, **kwargs):
             span.set_attribute("http.method", method.upper())
             span.set_attribute("http.url", url)
 
-            # TODO: Propagate the trace context via headers once we have a way
-            # to access propagators.
-
             headers = kwargs.setdefault("headers", {})
             propagators.inject(type(headers).__setitem__, headers)
             result = wrapped(self, method, url, *args, **kwargs)  # *** PROCEED
@@ -129,3 +127,10 @@ def disable():
     if getattr(Session.request, "opentelemetry_ext_requests_applied", False):
         original = Session.request.__wrapped__  # pylint:disable=no-member
         Session.request = original
+
+
+def disable_session(session):
+    """Disables instrumentation on the session object."""
+    if getattr(session.request, "opentelemetry_ext_requests_applied", False):
+        original = session.request.__wrapped__  # pylint:disable=no-member
+        session.request = types.MethodType(original, session)
diff --git a/ext/opentelemetry-ext-http-requests/tests/test_requests_integration.py b/ext/opentelemetry-ext-http-requests/tests/test_requests_integration.py
@@ -13,95 +13,50 @@
 # limitations under the License.
 
 import sys
-import unittest
-from unittest import mock
 
-import pkg_resources
+import httpretty
 import requests
 import urllib3
 
 import opentelemetry.ext.http_requests
-from opentelemetry import trace
+from opentelemetry import context, propagators, trace
+from opentelemetry.test.mock_httptextformat import MockHTTPTextFormat
+from opentelemetry.test.test_base import TestBase
 
 
-class TestRequestsIntegration(unittest.TestCase):
+class TestRequestsIntegration(TestBase):
+    URL = "http://httpbin.org/status/200"
 
-    # TODO: Copy & paste from test_wsgi_middleware
     def setUp(self):
-        self.span_attrs = {}
-        self.tracer_provider = trace.DefaultTracerProvider()
-        self.tracer = trace.DefaultTracer()
-        self.get_tracer_patcher = mock.patch.object(
-            self.tracer_provider,
-            "get_tracer",
-            autospec=True,
-            spec_set=True,
-            return_value=self.tracer,
-        )
-        self.get_tracer = self.get_tracer_patcher.start()
-        self.span_context_manager = mock.MagicMock()
-        self.span = mock.create_autospec(trace.Span, spec_set=True)
-        self.span.get_context.return_value = trace.INVALID_SPAN_CONTEXT
-        self.span_context_manager.__enter__.return_value = self.span
-
-        def setspanattr(key, value):
-            self.assertIsInstance(key, str)
-            self.span_attrs[key] = value
-
-        self.span.set_attribute = setspanattr
-        self.start_span_patcher = mock.patch.object(
-            self.tracer,
-            "start_as_current_span",
-            autospec=True,
-            spec_set=True,
-            return_value=self.span_context_manager,
-        )
-
-        mocked_response = requests.models.Response()
-        mocked_response.status_code = 200
-        mocked_response.reason = "Roger that!"
-        self.send_patcher = mock.patch.object(
-            requests.Session,
-            "send",
-            autospec=True,
-            spec_set=True,
-            return_value=mocked_response,
-        )
-
-        self.start_as_current_span = self.start_span_patcher.start()
-        self.send = self.send_patcher.start()
-
+        super().setUp()
         opentelemetry.ext.http_requests.enable(self.tracer_provider)
-        distver = pkg_resources.get_distribution(
-            "opentelemetry-ext-http-requests"
-        ).version
-        self.get_tracer.assert_called_with(
-            opentelemetry.ext.http_requests.__name__, distver
+        httpretty.enable()
+        httpretty.register_uri(
+            httpretty.GET, self.URL, body="Hello!",
         )
 
     def tearDown(self):
+        super().tearDown()
         opentelemetry.ext.http_requests.disable()
-        self.get_tracer_patcher.stop()
-        self.send_patcher.stop()
-        self.start_span_patcher.stop()
+        httpretty.disable()
 
     def test_basic(self):
-        url = "https://www.example.org/foo/bar?x=y#top"
-        requests.get(url=url)
-        self.assertEqual(1, len(self.send.call_args_list))
-        self.tracer.start_as_current_span.assert_called_with(  # pylint:disable=no-member
-            "/foo/bar", kind=trace.SpanKind.CLIENT
-        )
-        self.span_context_manager.__enter__.assert_called_with()
-        self.span_context_manager.__exit__.assert_called_with(None, None, None)
+        result = requests.get(self.URL)
+        self.assertEqual(result.text, "Hello!")
+        span_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(span_list), 1)
+        span = span_list[0]
+        self.assertIs(span.kind, trace.SpanKind.CLIENT)
+        self.assertEqual(span.name, "/status/200")
+
         self.assertEqual(
-            self.span_attrs,
+            span.attributes,
             {
                 "component": "http",
                 "http.method": "GET",
-                "http.url": url,
+                "http.url": self.URL,
                 "http.status_code": 200,
-                "http.status_text": "Roger that!",
+                "http.status_text": "OK",
             },
         )
 
@@ -114,20 +69,84 @@ def test_invalid_url(self):
             exception_type = ValueError
 
         with self.assertRaises(exception_type):
-            requests.post(url=url)
-        call_args = (
-            self.tracer.start_as_current_span.call_args  # pylint:disable=no-member
-        )
-        self.assertTrue(
-            call_args[0][0].startswith("<Unparsable URL"),
-            msg=self.tracer.start_as_current_span.call_args,  # pylint:disable=no-member
-        )
-        self.span_context_manager.__enter__.assert_called_with()
-        exitspan = self.span_context_manager.__exit__
-        self.assertEqual(1, len(exitspan.call_args_list))
-        self.assertIs(exception_type, exitspan.call_args[0][0])
-        self.assertIsInstance(exitspan.call_args[0][1], exception_type)
+            requests.post(url)
+
+        span_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(span_list), 1)
+        span = span_list[0]
+
+        self.assertTrue(span.name.startswith("<Unparsable URL"))
         self.assertEqual(
-            self.span_attrs,
+            span.attributes,
             {"component": "http", "http.method": "POST", "http.url": url},
         )
+
+    def test_disable(self):
+        opentelemetry.ext.http_requests.disable()
+        result = requests.get(self.URL)
+        self.assertEqual(result.text, "Hello!")
+        span_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(span_list), 0)
+
+    def test_disable_session(self):
+        session1 = requests.Session()
+        opentelemetry.ext.http_requests.disable_session(session1)
+
+        result = session1.get(self.URL)
+        self.assertEqual(result.text, "Hello!")
+        span_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(span_list), 0)
+
+        # Test that other sessions as well as global requests is still
+        # instrumented
+        session2 = requests.Session()
+        result = session2.get(self.URL)
+        self.assertEqual(result.text, "Hello!")
+        span_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(span_list), 1)
+
+        self.memory_exporter.clear()
+
+        result = requests.get(self.URL)
+        self.assertEqual(result.text, "Hello!")
+        span_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(span_list), 1)
+
+    def test_suppress_instrumentation(self):
+        token = context.attach(
+            context.set_value("suppress_instrumentation", True)
+        )
+        try:
+            result = requests.get(self.URL)
+            self.assertEqual(result.text, "Hello!")
+        finally:
+            context.detach(token)
+
+        span_list = self.memory_exporter.get_finished_spans()
+        self.assertEqual(len(span_list), 0)
+
+    def test_distributed_context(self):
+        previous_propagator = propagators.get_global_httptextformat()
+        try:
+            propagators.set_global_httptextformat(MockHTTPTextFormat())
+            result = requests.get(self.URL)
+            self.assertEqual(result.text, "Hello!")
+
+            span_list = self.memory_exporter.get_finished_spans()
+            self.assertEqual(len(span_list), 1)
+            span = span_list[0]
+
+            headers = dict(httpretty.last_request().headers)
+            self.assertIn(MockHTTPTextFormat.TRACE_ID_KEY, headers)
+            self.assertEqual(
+                str(span.get_context().trace_id),
+                headers[MockHTTPTextFormat.TRACE_ID_KEY],
+            )
+            self.assertIn(MockHTTPTextFormat.SPAN_ID_KEY, headers)
+            self.assertEqual(
+                str(span.get_context().span_id),
+                headers[MockHTTPTextFormat.SPAN_ID_KEY],
+            )
+
+        finally:
+            propagators.set_global_httptextformat(previous_propagator)
diff --git a/ext/opentelemetry-ext-opentracing-shim/setup.cfg b/ext/opentelemetry-ext-opentracing-shim/setup.cfg
@@ -44,5 +44,9 @@ install_requires =
     opentracing ~= 2.0
     opentelemetry-api == 0.7.dev0
 
+[options.extras_require]
+test =
+    opentelemetry-test == 0.7.dev0
+
 [options.packages.find]
 where = src
diff --git a/ext/opentelemetry-ext-opentracing-shim/tests/test_shim.py b/ext/opentelemetry-ext-opentracing-shim/tests/test_shim.py
@@ -16,26 +16,15 @@
 # pylint:disable=no-member
 
 import time
-import typing
 from unittest import TestCase
 
 import opentracing
 
 import opentelemetry.ext.opentracing_shim as opentracingshim
 from opentelemetry import propagators, trace
-from opentelemetry.context import Context
 from opentelemetry.ext.opentracing_shim import util
 from opentelemetry.sdk.trace import TracerProvider
-from opentelemetry.trace.propagation import (
-    get_span_from_context,
-    set_span_in_context,
-)
-from opentelemetry.trace.propagation.httptextformat import (
-    Getter,
-    HTTPTextFormat,
-    HTTPTextFormatT,
-    Setter,
-)
+from opentelemetry.test.mock_httptextformat import MockHTTPTextFormat
 
 
 class TestShim(TestCase):
@@ -542,46 +531,3 @@ def test_extract_binary(self):
         # Verify exception for non supported binary format.
         with self.assertRaises(opentracing.UnsupportedFormatException):
             self.shim.extract(opentracing.Format.BINARY, bytearray())
-
-
-class MockHTTPTextFormat(HTTPTextFormat):
-    """Mock propagator for testing purposes."""
-
-    TRACE_ID_KEY = "mock-traceid"
-    SPAN_ID_KEY = "mock-spanid"
-
-    def extract(
-        self,
-        get_from_carrier: Getter[HTTPTextFormatT],
-        carrier: HTTPTextFormatT,
-        context: typing.Optional[Context] = None,
-    ) -> Context:
-        trace_id_list = get_from_carrier(carrier, self.TRACE_ID_KEY)
-        span_id_list = get_from_carrier(carrier, self.SPAN_ID_KEY)
-
-        if not trace_id_list or not span_id_list:
-            return set_span_in_context(trace.INVALID_SPAN)
-
-        return set_span_in_context(
-            trace.DefaultSpan(
-                trace.SpanContext(
-                    trace_id=int(trace_id_list[0]),
-                    span_id=int(span_id_list[0]),
-                    is_remote=True,
-                )
-            )
-        )
-
-    def inject(
-        self,
-        set_in_carrier: Setter[HTTPTextFormatT],
-        carrier: HTTPTextFormatT,
-        context: typing.Optional[Context] = None,
-    ) -> None:
-        span = get_span_from_context(context)
-        set_in_carrier(
-            carrier, self.TRACE_ID_KEY, str(span.get_context().trace_id)
-        )
-        set_in_carrier(
-            carrier, self.SPAN_ID_KEY, str(span.get_context().span_id)
-        )
diff --git a/tests/util/src/opentelemetry/test/mock_httptextformat.py b/tests/util/src/opentelemetry/test/mock_httptextformat.py
diff --git a/tests/util/src/opentelemetry/test/test_base.py b/tests/util/src/opentelemetry/test/test_base.py
diff --git a/tox.ini b/tox.ini
