Propagate tracing headers for requests module (#338)

This includes the header generated by current trace propagator's method
when wrapping requests.Session. It fixes the cross request propagation
for requests module.

The fix targets wrap on requests.Session object, because it covers all common
requests module use cases:

- Calling requests.{get,post,put,...} instantiates requests.Session,
so the header would end up properly applied from the Session instance.

- Instantiating requests.Session for caching purposes would directly
apply the header.

Author: irekatroszko

opencensus/trace/ext/requests/trace.py

@@ -82,6 +82,12 @@ def wrap_session_request(wrapped, instance, args, kwargs):
     _span.name = '[requests]{}'.format(method)
     _span.span_kind = span_module.SpanKind.CLIENT
 
+    tracer_headers = _tracer.propagator.to_headers(
+        _tracer.span_context)
+
+    kwargs.setdefault('headers', {}).update(
+        tracer_headers)
+
     # Add the requests url to attributes
     _tracer.add_attribute_to_current_span(HTTP_URL, url)
 

tests/unit/trace/ext/requests/test_requests_trace.py

@@ -78,12 +78,11 @@ def test_wrap_requests(self):
         self.assertEqual(expected_name, mock_tracer.current_span.name)
 
     def test_wrap_session_request(self):
-        def wrapped(*args, **kwargs):
-            result = mock.Mock()
-            result.status_code = 200
-            return result
+        wrapped = mock.Mock(return_value=mock.Mock(status_code=200))
 
-        mock_tracer = MockTracer()
+        mock_tracer = MockTracer(propagator=mock.Mock(
+            to_headers= lambda x: {'x-trace': 'some-value'})
+        )
 
         patch = mock.patch(
             'opencensus.trace.ext.requests.trace.execution_context.'
@@ -92,10 +91,11 @@ def wrapped(*args, **kwargs):
 
         url = 'http://localhost:8080'
         request_method = 'POST'
+        kwargs = {}
 
         with patch:
             result = trace.wrap_session_request(
-                wrapped, 'Session.request', (request_method, url), {})
+                wrapped, 'Session.request', (request_method, url), kwargs)
 
         expected_attributes = {
             'http.url': url,
@@ -106,12 +106,79 @@ def wrapped(*args, **kwargs):
                          mock_tracer.current_span.span_kind)
         self.assertEqual(expected_attributes,
                          mock_tracer.current_span.attributes)
+        self.assertEqual(kwargs['headers']['x-trace'], 'some-value')
         self.assertEqual(expected_name, mock_tracer.current_span.name)
 
+    def test_header_is_passed_in(self):
+        wrapped = mock.Mock(return_value=mock.Mock(status_code=200))
+        mock_tracer = MockTracer(propagator=mock.Mock(
+            to_headers= lambda x: {'x-trace': 'some-value'})
+        )
+        
+        patch = mock.patch(
+            'opencensus.trace.ext.requests.trace.execution_context.'
+            'get_opencensus_tracer',
+            return_value=mock_tracer)
+
+        url = 'http://localhost:8080'
+        request_method = 'POST'
+        kwargs = {}
+
+        with patch:
+            result = trace.wrap_session_request(
+                wrapped, 'Session.request', (request_method, url), kwargs)
+
+        self.assertEqual(kwargs['headers']['x-trace'], 'some-value')
+
+    def test_headers_are_preserved(self):
+        wrapped = mock.Mock(return_value=mock.Mock(status_code=200))
+        mock_tracer = MockTracer(propagator=mock.Mock(
+            to_headers= lambda x: {'x-trace': 'some-value'})
+        )
+
+        patch = mock.patch(
+            'opencensus.trace.ext.requests.trace.execution_context.'
+            'get_opencensus_tracer',
+            return_value=mock_tracer)
+
+        url = 'http://localhost:8080'
+        request_method = 'POST'
+        kwargs = {'headers': {'key': 'value'}}
+
+        with patch:
+            result = trace.wrap_session_request(
+                wrapped, 'Session.request', (request_method, url), kwargs)
+
+        self.assertEqual(kwargs['headers']['key'], 'value')
+        self.assertEqual(kwargs['headers']['x-trace'], 'some-value')
+
+
+    def test_tracer_headers_are_overwritten(self):
+        wrapped = mock.Mock(return_value=mock.Mock(status_code=200))
+        mock_tracer = MockTracer(propagator=mock.Mock(
+            to_headers= lambda x: {'x-trace': 'some-value'})
+        )
+
+        patch = mock.patch(
+            'opencensus.trace.ext.requests.trace.execution_context.'
+            'get_opencensus_tracer',
+            return_value=mock_tracer)
+
+        url = 'http://localhost:8080'
+        request_method = 'POST'
+        kwargs = {'headers': {'x-trace': 'original-value'}}
+
+        with patch:
+            result = trace.wrap_session_request(
+                wrapped, 'Session.request', (request_method, url), kwargs)
+
+        self.assertEqual(kwargs['headers']['x-trace'], 'some-value')
 
 class MockTracer(object):
-    def __init__(self):
+    def __init__(self, propagator=None):
         self.current_span = None
+        self.span_context = {}
+        self.propagator = propagator
 
     def start_span(self):
         span = mock.Mock()