feat(urrlib3): collecting payload (#356)

epsagon/events/urllib3.py

@@ -76,7 +76,7 @@ def __init__(self, wrapped, instance, args, kwargs, start_time, response,
             add_data_if_needed(
                 self.resource['metadata'],
                 'request_headers',
-                dict(headers)
+                dict(headers) if headers else {}
             )
 
             add_data_if_needed(
@@ -98,7 +98,7 @@ def update_response(self, response):
         :return: None
         """
         self.resource['metadata']['status_code'] = response.status
-        headers = dict(response.getheaders())
+        headers = dict(response.getheaders()) if response.getheaders() else {}
         self.resource = update_http_headers(
             self.resource,
             headers
@@ -113,7 +113,13 @@ def update_response(self, response):
                 'response_headers',
                 headers
             )
-            response_body = response.peek()
+            response_body = (
+                response.peek()
+                if getattr(response, 'peek', None)
+                else response.data
+                if response.tell() > 0
+                else None
+            )
             if isinstance(response_body, bytes):
                 try:
                     response_body = response_body.decode('utf-8')

epsagon/modules/urllib.py

@@ -3,7 +3,7 @@
 """
 
 from __future__ import absolute_import
-import wrapt
+from epsagon.utils import patch_once
 from epsagon.modules.general_wrapper import wrapper
 from ..events.urllib import UrllibEventFactory
 
@@ -27,7 +27,7 @@ def patch():
     """
 
     try:
-        wrapt.wrap_function_wrapper(
+        patch_once(
             'urllib.request',
             'OpenerDirector._open',
             _wrapper

epsagon/modules/urllib3.py

@@ -4,7 +4,7 @@
 
 from __future__ import absolute_import
 import uuid
-import wrapt
+from epsagon.utils import patch_once
 from epsagon.modules.general_wrapper import wrapper
 from ..events.urllib3 import Urllib3EventFactory
 from ..http_filters import is_blacklisted_url
@@ -82,7 +82,7 @@ def patch():
     """
 
     try:
-        wrapt.wrap_function_wrapper(
+        patch_once(
             'urllib3',
             'HTTPConnectionPool.urlopen',
             _wrapper

tests/modules/test_urllib3.py

@@ -0,0 +1,95 @@
+import pytest
+import urllib3
+import epsagon.wrappers.python_function
+from epsagon.trace import trace_factory
+import epsagon.runners.python_function
+import epsagon.constants
+
+
+TEST_DOMAIN = 'jsonplaceholder.typicode.com'
+TEST_PATH = '/todos/1'
+
+def setup_function(func):
+    trace_factory.use_single_trace = True
+    epsagon.constants.COLD_START = True
+
+def test_no_data_capture_with_urlopen(trace_transport):
+    def use_urlopen():
+        headers = {
+            'Content-Type': 'application/json'
+        }
+
+        http = urllib3.PoolManager()
+        conn = http.connection_from_url(TEST_DOMAIN)
+        urllib_response = conn.urlopen(
+            method='GET',
+            url=TEST_PATH,
+            body='',
+            headers=headers,
+            preload_content=False,
+            decode_content=False,
+        )
+        return urllib_response
+
+    response = use_urlopen()
+    data = response.read()
+
+    @epsagon.wrappers.python_function.python_wrapper
+    def wrapped_function():
+        return use_urlopen()
+
+    wrapped_response = wrapped_function()
+    wrapped_data = wrapped_response.read()
+    # We expect in this case that the data will be available in the buffer
+    assert(len(wrapped_data) > 0)
+    assert(len(wrapped_response.data) == 0)
+    urllib3_event = trace_transport.last_trace.events[-1]
+    # Payload should not be collected in the case of raw-stream usage.
+    assert(urllib3_event.resource['metadata']['response_body'] is None)
+    # Compare un-instrumented vs instrumented data
+    assert (len(data) > 0)
+    assert (len(response.data) == 0)
+    assert(wrapped_data == data)
+    assert(wrapped_response.data == response.data)
+
+@pytest.mark.parametrize("preload_content", [True, False])
+def test_data_capture_with_pool_manager(preload_content, trace_transport):
+    def use_poolmanager():
+        headers = {
+            'Content-Type': 'application/json'
+        }
+
+        http = urllib3.PoolManager()
+        return http.request(
+            'GET', TEST_DOMAIN + TEST_PATH,
+            headers=headers, body='',
+            preload_content=preload_content
+        )
+
+    response = use_poolmanager()
+    data = response.read()
+
+    @epsagon.wrappers.python_function.python_wrapper
+    def wrapped_function():
+        return use_poolmanager()
+
+    wrapped_response = wrapped_function()
+    wrapped_data = wrapped_response.read()
+    if not preload_content:
+        # In this case data will not be available in the buffer
+        assert (len(wrapped_data) > 0)
+        # But, data will not be stored in the `data` object
+        assert (len(wrapped_response.data) == 0)
+        # Compare un-instrumented vs instrumented data
+        assert (len(data) > 0)
+        assert (len(response.data) == 0)
+    else:
+        # In this case data will not be available in the buffer
+        assert (len(wrapped_data) == 0)
+        # But, data will be stored in the `data` object
+        assert (len(wrapped_response.data) > 0)
+        # Compare un-instrumented vs instrumented data
+        assert (len(data) == 0)
+        assert (len(response.data) > 0)
+    assert(wrapped_data == data)
+    assert(wrapped_response.data == response.data)

tests/test_utils.py

@@ -1,3 +1,4 @@
+import pytest
 import epsagon.trace
 import epsagon.utils
 import epsagon.http_filters
@@ -7,8 +8,13 @@
 def setup_function(func):
     trace_factory.get_or_create_trace()
 
+@pytest.yield_fixture
+def blacklist_urls_state_keeper():
+    original = epsagon.http_filters.BLACKLIST_URLS  # Storing old state
+    yield
+    epsagon.http_filters.BLACKLIST_URLS = original  # Restoring old state
 
-def test_blacklist_url():
+def test_blacklist_url(blacklist_urls_state_keeper):
     """
     Test is_blacklisted_url functionality.
     :return: None
@@ -36,7 +42,6 @@ def test_original_blacklist_url():
     Validate original needed URLs are in.
     :return: None
     """
-
     assert epsagon.http_filters.is_blacklisted_url('http://tc.us-east-1.epsagon.com')
     assert epsagon.http_filters.is_blacklisted_url('https://client.tc.epsagon.com')