add custom headers to before-call and access them in request-created

Signed-off-by: Varsha GS <[email protected]>

instana/instrumentation/boto3_inst.py

@@ -46,20 +46,20 @@ def lambda_inject_context(payload, scope):
             logger.debug("non-fatal lambda_inject_context: ", exc_info=True)
 
 
-    @wrapt.patch_function_wrapper("botocore.hooks", "HierarchicalEmitter.emit_until_response")
-    def emit_until_response_with_instana(wrapped, instance, args, kwargs):
+    @wrapt.patch_function_wrapper("botocore.hooks", "HierarchicalEmitter.emit")
+    def emit_request_created_with_instana(wrapped, instance, args, kwargs):
         active_tracer = get_active_tracer()
 
-        # If we're not tracing or the event emitted is not before-call, just return;
-        if active_tracer is None or args[0].split(".")[0] != "before-call":
+        # If we're not tracing or the event emitted is not request-created, just return;
+        if active_tracer is None or args[0].split(".")[0] != "request-created":
             return wrapped(*args, **kwargs)
 
         span = active_tracer.active_span
-        if "custom_request_headers" in kwargs["context"]:
-            extract_custom_headers(span, kwargs["context"]["custom_request_headers"])
+        extract_custom_headers(span, kwargs["request"].headers)
 
         return wrapped(*args, **kwargs)
 
+
     @wrapt.patch_function_wrapper('botocore.client', 'BaseClient._make_api_call')
     def make_api_call_with_instana(wrapped, instance, arg_list, kwargs):
         # pylint: disable=protected-access

tests/clients/boto3/test_boto3_s3.py

@@ -22,8 +22,8 @@
 
 
 class TestS3(unittest.TestCase):
-    def aws_credentials(self):
-        """Mocked AWS Credentials for moto."""
+    def set_aws_credentials(self):
+        """ Mocked AWS Credentials for moto """
         os.environ['AWS_ACCESS_KEY_ID'] = 'testing'
         os.environ['AWS_SECRET_ACCESS_KEY'] = 'testing'
         os.environ['AWS_SECURITY_TOKEN'] = 'testing'
@@ -33,16 +33,25 @@ def setUp(self):
         """ Clear all spans before a test run """
         self.recorder = tracer.recorder
         self.recorder.clear_spans()
-        self.aws_credentials()
+        self.set_aws_credentials()
         self.mock = mock_aws()
         self.mock.start()
         self.s3 = boto3.client('s3', region_name='us-east-1')
 
+    def unset_aws_credentials(self):
+        """ Reset all environment variables of consequence """
+        variable_names = (
+                "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY",
+                "AWS_SECURITY_TOKEN", "AWS_SESSION_TOKEN"
+                )
+
+        for variable_name in variable_names:
+            os.environ.pop(variable_name, None)
 
     def tearDown(self):
         # Stop Moto after each test
         self.mock.stop()
-
+        self.unset_aws_credentials()
 
     def test_vanilla_create_bucket(self):
         self.s3.create_bucket(Bucket="aws_bucket_name")
@@ -88,9 +97,8 @@ def test_s3_create_bucket(self):
 
     def test_s3_list_buckets(self):
         with tracer.start_active_span('test'):
-            self.s3.list_buckets()
+            result = self.s3.list_buckets()
 
-        result = self.s3.list_buckets()
         self.assertEqual(0, len(result['Buckets']))
         self.assertEqual(result['ResponseMetadata']['HTTPStatusCode'], 200)
 
@@ -283,15 +291,13 @@ def test_request_header_capture(self):
                 'X-Capture-This': 'this',
                 'X-Capture-That': 'that'
             }
-
-        # We set the custom headers in the request context instead of params 
-        # because later in the processing of the request, there is a parameter validation step, 
-        # which doesn't allow for custom arguments. 
-        def process_custom_arguments(params, context, **kwargs):
-            if "custom_request_headers" not in context:
-                context["custom_request_headers"] = request_headers
-
-        event_system.register('before-parameter-build.s3.CreateBucket', process_custom_arguments)
+
+        # Create a function that adds custom headers
+        def add_custom_header_before_call(params, **kwargs):
+            params['headers'].update(request_headers)
+
+        # Register the function to an event.
+        event_system.register('before-call.s3.CreateBucket', add_custom_header_before_call)
 
         with tracer.start_active_span('test'):
             result = self.s3.create_bucket(Bucket="aws_bucket_name")
@@ -354,7 +360,7 @@ def modify_after_call_args(parsed, **kwargs):
         event_system.register('after-call.s3.CreateBucket', modify_after_call_args)
 
         with tracer.start_active_span('test'):
-            result = self.s3.create_bucket(Bucket="aws_bucket_name")
+            self.s3.create_bucket(Bucket="aws_bucket_name")
 
         result = self.s3.list_buckets()
         self.assertEqual(1, len(result['Buckets']))