From 28e97dd128a23e5e36ecee558240bb676bb7931e Mon Sep 17 00:00:00 2001
From: w- <w-@users.noreply.github.com>
Date: Wed, 31 May 2017 19:22:19 +0700
Subject: [PATCH] Add ability to reset request headers on client attribute

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.
---
 sendgrid/sendgrid.py  | 13 ++++++++++---
 test/test_sendgrid.py | 31 +++++++++++++++++++++++++++++++
 2 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/sendgrid/sendgrid.py b/sendgrid/sendgrid.py
index 61c664d88..3a4184027 100644
--- 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
index 1f16dcf2b..11331fbab 100644
--- 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}