From ae9d1e45226abd641867830c2102828ce35f5568 Mon Sep 17 00:00:00 2001
From: Brandon Squizzato <bsquizza@redhat.com>
Date: Wed, 28 Nov 2018 14:17:48 -0500
Subject: [PATCH] Add option for setting a custom requests.Session

---
 pyswagger/contrib/client/requests.py          | 17 +++++++--
 .../tests/contrib/client/test_requests.py     | 38 +++++++++++++++++++
 2 files changed, 51 insertions(+), 4 deletions(-)

diff --git a/pyswagger/contrib/client/requests.py b/pyswagger/contrib/client/requests.py
index a7bce1d..2b45332 100644
--- a/pyswagger/contrib/client/requests.py
+++ b/pyswagger/contrib/client/requests.py
@@ -10,19 +10,29 @@ class Client(BaseClient):
 
     __schemes__ = set(['http', 'https'])
 
-    def __init__(self, auth=None, send_opt=None):
+    def set_session(self, session):
+        """ set the requests.Session used by the client
+
+        :param session requests.Session: instance of Session client should use
+        """
+        if not isinstance(session, Session):
+            raise TypeError("session is not an instance of requests.Session")
+        self.__s = session
+
+    def __init__(self, auth=None, send_opt=None, session=None):
         """ constructor
 
         :param auth pyswagger.SwaggerAuth: auth info used when requesting
         :param send_opt dict: options used in requests.send, ex verify=False
+        :param session requests.Session: optional, set the client Session
         """
         super(Client, self).__init__(auth)
         if send_opt is None:
             send_opt = {}
-
-        self.__s = Session()
         self.__send_opt = send_opt
 
+        self.set_session(session or Session())
+
     def request(self, req_and_resp, opt=None, headers=None):
         """
         """
@@ -75,4 +85,3 @@ def append(name, obj):
         )
 
         return resp
-
diff --git a/pyswagger/tests/contrib/client/test_requests.py b/pyswagger/tests/contrib/client/test_requests.py
index 004a1ab..bf2285c 100644
--- a/pyswagger/tests/contrib/client/test_requests.py
+++ b/pyswagger/tests/contrib/client/test_requests.py
@@ -8,6 +8,7 @@
 import json
 import six
 import os
+import requests
 
 
 app = App._create_(get_test_data_folder(version='1.2', which='wordnik'))
@@ -264,3 +265,40 @@ def test_custom_headers_multiple_values_to_one_key(self):
         )
         self.assertEqual(httpretty.last_request().headers['X-TEST-HEADER'], 'bbb')
 
+
+@httpretty.activate
+class SessionTestCase(unittest.TestCase):
+    """ test case for custom session """
+
+    def test_init_with_session(self):
+        session = requests.Session()
+        session.headers = {'X-TEST-HEADER': 'aaa'}
+        session_client = Client(session=session)
+
+        httpretty.register_uri(
+            httpretty.PUT,
+            'http://petstore.swagger.wordnik.com/api/pet',
+            status=200
+        )
+
+        session_client.request(app.op['updatePet'](body=pet_QQQ))
+
+        self.assertEqual(httpretty.last_request().headers['X-TEST-HEADER'], 'aaa')
+
+    def test_set_session(self):
+        session = requests.Session()
+        session.headers = {'X-TEST-HEADER': 'aaa'}
+        default_client = Client()
+        session_client = Client(session=session)
+
+        httpretty.register_uri(
+            httpretty.PUT,
+            'http://petstore.swagger.wordnik.com/api/pet',
+            status=200
+        )
+
+        session_client.request(app.op['updatePet'](body=pet_QQQ))
+        self.assertEqual(httpretty.last_request().headers['X-TEST-HEADER'], 'aaa')
+
+        default_client.request(app.op['updatePet'](body=pet_QQQ))
+        self.assertFalse('X-TEST-HEADER' in httpretty.last_request().headers)