From 7e441e5c5e670922b37cb63023494f4988f49657 Mon Sep 17 00:00:00 2001 From: Alexandre Yang Date: Fri, 14 Feb 2020 19:04:32 +0100 Subject: [PATCH 1/2] Add extra_headers option to http method call --- .../datadog_checks/base/utils/http.py | 10 ++++++++-- datadog_checks_base/tests/test_http.py | 19 +++++++++++++++++++ 2 files changed, 27 insertions(+), 2 deletions(-) diff --git a/datadog_checks_base/datadog_checks/base/utils/http.py b/datadog_checks_base/datadog_checks/base/utils/http.py index 8cb8b6c5a493a..b18f8cabaa817 100644 --- a/datadog_checks_base/datadog_checks/base/utils/http.py +++ b/datadog_checks_base/datadog_checks/base/utils/http.py @@ -302,14 +302,20 @@ def _request(self, method, url, options): if persist is None: persist = self.persist_connections + new_options = self.populate_options(options) + + extra_headers = options.pop('extra_headers', None) + if extra_headers is not None: + new_options['headers'].update(extra_headers) + with ExitStack() as stack: for hook in self.request_hooks: stack.enter_context(hook()) if persist: - return getattr(self.session, method)(url, **self.populate_options(options)) + return getattr(self.session, method)(url, **new_options) else: - return getattr(requests, method)(url, **self.populate_options(options)) + return getattr(requests, method)(url, **new_options) def populate_options(self, options): # Avoid needless dictionary update if there are no options diff --git a/datadog_checks_base/tests/test_http.py b/datadog_checks_base/tests/test_http.py index 88900642d4e99..d5848ac20ebcc 100644 --- a/datadog_checks_base/tests/test_http.py +++ b/datadog_checks_base/tests/test_http.py @@ -119,6 +119,25 @@ def test_config_extra_headers_string_values(self): assert http.options['headers'] == {'User-Agent': 'Datadog Agent/0.0.0', 'answer': '42'} + def test_extra_headers_on_http_method_call(self): + instance = {'extra_headers': {'answer': 42}} + init_config = {} + http = RequestsWrapper(instance, init_config) + + with mock.patch("requests.get") as get: + http.get("http://example.com/hello", extra_headers={"foo": "bar"}) + + expected_options = {'foo': 'bar', 'User-Agent': 'Datadog Agent/0.0.0', 'answer': '42'} + get.assert_called_with( + "http://example.com/hello", + headers=expected_options, + auth=None, + cert=None, + proxies=None, + timeout=(10.0, 10.0), + verify=True, + ) + class TestVerify: def test_config_default(self): From 166fdf7efd54d509e8df161cee54e41aa909b7cd Mon Sep 17 00:00:00 2001 From: Alexandre Yang Date: Fri, 14 Feb 2020 22:12:17 +0100 Subject: [PATCH 2/2] Make copy --- datadog_checks_base/datadog_checks/base/utils/http.py | 1 + datadog_checks_base/tests/test_http.py | 7 ++++++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/datadog_checks_base/datadog_checks/base/utils/http.py b/datadog_checks_base/datadog_checks/base/utils/http.py index b18f8cabaa817..5e636991f6511 100644 --- a/datadog_checks_base/datadog_checks/base/utils/http.py +++ b/datadog_checks_base/datadog_checks/base/utils/http.py @@ -306,6 +306,7 @@ def _request(self, method, url, options): extra_headers = options.pop('extra_headers', None) if extra_headers is not None: + new_options['headers'] = new_options['headers'].copy() new_options['headers'].update(extra_headers) with ExitStack() as stack: diff --git a/datadog_checks_base/tests/test_http.py b/datadog_checks_base/tests/test_http.py index d5848ac20ebcc..d3470cfb9e944 100644 --- a/datadog_checks_base/tests/test_http.py +++ b/datadog_checks_base/tests/test_http.py @@ -124,8 +124,9 @@ def test_extra_headers_on_http_method_call(self): init_config = {} http = RequestsWrapper(instance, init_config) + extra_headers = {"foo": "bar"} with mock.patch("requests.get") as get: - http.get("http://example.com/hello", extra_headers={"foo": "bar"}) + http.get("http://example.com/hello", extra_headers=extra_headers) expected_options = {'foo': 'bar', 'User-Agent': 'Datadog Agent/0.0.0', 'answer': '42'} get.assert_called_with( @@ -138,6 +139,10 @@ def test_extra_headers_on_http_method_call(self): verify=True, ) + # make sure the original headers are not modified + assert http.options['headers'] == {'User-Agent': 'Datadog Agent/0.0.0', 'answer': '42'} + assert extra_headers == {"foo": "bar"} + class TestVerify: def test_config_default(self):