Add ability to reset request headers on client attribute

w- · w- · commit 28e97dd128a2 · 2017-05-31T19:22:19.000+07:00
Because of the way SendGridAPIClient and python_http_client are implemented, request headers are always cached. This becomes a problem when executing a sequence of API requests that require varying request header values using the same SendGridAPIClient instance.
This method allows the caller to reset the request headers on the client instance to their default.
diff --git a/sendgrid/sendgrid.py b/sendgrid/sendgrid.py
@@ -26,7 +26,13 @@ def __init__(self, **opts):
         self.host = opts.get('host', 'https://api.sendgrid.com')
         self.version = __version__
 
+        headers = self._get_default_headers()
 
+        self.client = python_http_client.Client(host=self.host,
+                                                request_headers=headers,
+                                                version=3)
+
+    def _get_default_headers(self):
         headers = {
             "Authorization": 'Bearer {0}'.format(self._apikey),
             "User-agent": self.useragent,
@@ -35,9 +41,10 @@ def __init__(self, **opts):
         if self._impersonate_subuser:
             headers['On-Behalf-Of'] = self._impersonate_subuser
 
-        self.client = python_http_client.Client(host=self.host,
-                                                request_headers=headers,
-                                                version=3)
+        return headers
+
+    def reset_request_headers(self):
+        self.client.request_headers = self._get_default_headers()
 
     @property
     def apikey(self):
diff --git a/test/test_sendgrid.py b/test/test_sendgrid.py
@@ -78,6 +78,37 @@ def test_useragent(self):
     def test_host(self):
         self.assertEqual(self.sg.host, self.host)
 
+    def test_get_default_headers(self):
+        headers = self.sg._get_default_headers()
+        self.assertIn('Authorization', headers)
+        self.assertIn('User-agent', headers)
+        self.assertIn('Accept', headers)
+        self.assertNotIn('On-Behalf-Of', headers)
+
+        self.sg._impersonate_subuser = 'ladida@testsubuser.sendgrid'
+        headers = self.sg._get_default_headers()
+        self.assertIn('Authorization', headers)
+        self.assertIn('User-agent', headers)
+        self.assertIn('Accept', headers)
+        self.assertIn('On-Behalf-Of', headers)
+
+    def test_reset_request_headers(self):
+        addl_headers = {
+            'blah': 'test value',
+            'blah2x': 'another test value',
+        }
+        self.sg.client.request_headers.update(addl_headers)
+        self.assertIn('blah', self.sg.client.request_headers)
+        self.assertIn('blah2x', self.sg.client.request_headers)
+
+        self.sg.reset_request_headers()
+        self.assertNotIn('blah', self.sg.client.request_headers)
+        self.assertNotIn('blah2x', self.sg.client.request_headers)
+
+        for k,v in self.sg._get_default_headers().items():
+            self.assertEqual(v, self.sg.client.request_headers[k])
+
+
     def test_access_settings_activity_get(self):
         params = {'limit': 1}
         headers = {'X-Mock': 200}
