updated IMS class to use session. RE DBOARD3-1095

inventory_provider/db/ims.py

@@ -199,10 +199,23 @@ class IMS(object):
         self.base_url = base_url
         self.username = username
         self.password = password
-        self.verify_ssl = verify_ssl
-        self.bearer_token = None
-        self.bearer_token_init_time = 0
         self.reconnect_attempts = 0
+        self.session = requests.session()
+        self.session.verify = verify_ssl
+
+    @property
+    def bearer_token(self):
+        try:
+            return self.session.headers['Authorization'].replace('Bearer ', '')
+        except KeyError:
+            return None
+
+    @bearer_token.setter
+    def bearer_token(self, value):
+        if value is None:
+            self.session.headers.pop('Authorization', None)
+        else:
+            self.session.headers.update({'Authorization': f'Bearer {value}'})
 
     def __enter__(self):
         self._init_bearer_token()
@@ -220,10 +233,7 @@ class IMS(object):
 
         for i in range(IMS.PERMITTED_RECONNECT_ATTEMPTS):
             try:
-                response = requests.get(
-                    self.base_url + self.HEARTBEAT_PATH,
-                    headers={'Authorization': f'Bearer {self.bearer_token}'},
-                    verify=self.verify_ssl)
+                response = self.session.get(self.base_url + self.HEARTBEAT_PATH)
                 return response.status_code == requests.codes.ok
             except HTTPError as e:
                 logger.error(f'Error checking heartbeat (attempt {i + 1}): {e}')
@@ -241,9 +251,9 @@ class IMS(object):
 
         logger.debug(f'Logging in - Username: {self.username}'
                      f' - URL: {self.base_url + self.LOGIN_PATH}')
-        response = requests.post(
+        response = self.session.post(
             self.base_url + self.LOGIN_PATH,
-            auth=(self.username, self.password), verify=self.verify_ssl)
+            auth=(self.username, self.password))
         response.raise_for_status()
         self.bearer_token = response.text
 
@@ -251,11 +261,7 @@ class IMS(object):
         if self.bearer_token is None:
             return
         logout_url = self.base_url + self.LOGOUT_PATH
-        response = requests.get(
-            logout_url,
-            headers={'Authorization': f'Bearer {self.bearer_token}'},
-            verify=self.verify_ssl
-        )
+        response = self.session.get(logout_url)
         try:
             response.raise_for_status()
         except HTTPError as e:
@@ -267,10 +273,9 @@ class IMS(object):
             self._init_bearer_token()
         while True:
             logger.info('Clearing Dynamic Context Cache')
-            response = requests.put(
+            response = self.session.put(
                 f'{self.base_url + IMS.IMS_PATH}/ClearDynamicContextCache',
-                headers={'Authorization': f'Bearer {self.bearer_token}'},
-                verify=self.verify_ssl)
+            )
             if response.status_code == 401:
                 self._init_bearer_token()
                 continue
@@ -332,10 +337,7 @@ class IMS(object):
             return source
 
         while True:
-            response = requests.get(
-                url,
-                headers={'Authorization': f'Bearer {self.bearer_token}'},
-                params=params, verify=self.verify_ssl)
+            response = self.session.get(url, params=params)
             if _is_invalid_login_state(response):
                 self._init_bearer_token()
             else:

test/test_ims.py

@@ -22,16 +22,16 @@ PASSWORD = "dummy_password"
 
 @pytest.fixture
 def ims_instance(mocker):
-    mock_post = mocker.patch('inventory_provider.db.ims.requests.post')
+    ims_instance = PatchedIMS(BASE_URL, USERNAME, PASSWORD)
+    mock_post = mocker.patch.object(ims_instance.session, 'post')
     mock_post.return_value.status_code = 200
     mock_post.return_value.text = 'dummy_bt'
-    return PatchedIMS(BASE_URL, USERNAME, PASSWORD)
+    return ims_instance
 
 
 def test_init_bearer_token(ims_instance):
 
     ims_instance._init_bearer_token()
-
     assert ims_instance.bearer_token == 'dummy_bt'
     assert ims_instance.reconnect_attempts == 0
 
@@ -44,7 +44,7 @@ def test_heartbeat(mocker, ims_instance):
     # set the bearer token
     ims_instance.bearer_token = 'dummy_bt'
 
-    mocked_get = mocker.patch('inventory_provider.db.ims.requests.get')
+    mocked_get = mocker.patch.object(ims_instance.session, 'get')
 
     mocked_get.return_value.status_code = 200
     assert ims_instance._is_heartbeat_ok()
@@ -76,7 +76,7 @@ def test_get_entity_by_id(mocker, ims_instance):
     entity_id = '12345'
     expected_entity = {'id': entity_id, 'name': 'Test Entity'}
 
-    mocked_get = mocker.patch('inventory_provider.db.ims.requests.get')
+    mocked_get = mocker.patch.object(ims_instance.session, 'get')
     mocked_get.return_value.status_code = 200
     mocked_get.return_value.json.return_value = expected_entity
 
@@ -90,7 +90,7 @@ def test_get_entity_by_name(mocker, ims_instance):
     entity_name = 'dummy_entity_name'
     expected_entity = {'id': '1234', 'name': entity_name}
 
-    mocked_get = mocker.patch('inventory_provider.db.ims.requests.get')
+    mocked_get = mocker.patch.object(ims_instance.session, 'get')
     mocked_get.return_value.status_code = 200
     mocked_get.return_value.json.return_value = expected_entity
 
@@ -104,7 +104,7 @@ def test_get_filtered_entities(mocker, ims_instance):
     entity_filter = 'dummy_param=dummy value'
     expected_result = [{'id': '1234', 'name': 'entity_name'}]
 
-    mocked_get = mocker.patch('inventory_provider.db.ims.requests.get')
+    mocked_get = mocker.patch.object(ims_instance.session, 'get')
     mocked_get.return_value.status_code = 200
     mocked_get.return_value.json.return_value = expected_result
 
@@ -117,7 +117,7 @@ def test_get_entities(mocker, ims_instance):
     entity = 'dummy_entity_type'
     expected_result = [{'id': '1234', 'name': 'entity_name'}]
 
-    mocked_get = mocker.patch('inventory_provider.db.ims.requests.get')
+    mocked_get = mocker.patch.object(ims_instance.session, 'get')
     mocked_get.return_value.status_code = 200
     mocked_get.return_value.json.return_value = expected_result
 
@@ -140,7 +140,7 @@ def test_get_paged_entities(step_count, ims_instance, mocker):
         {'id': '12348', 'name': 'entity_name_8'},
         {'id': '12349', 'name': 'entity_name_9'},
     ]
-    mocked_get = mocker.patch('inventory_provider.db.ims.requests.get')
+    mocked_get = mocker.patch.object(ims_instance.session, 'get')
     mocked_responses = []
 
     for i in range(0, len(expected_result) + 1, step_count):
@@ -159,10 +159,7 @@ def test_get_paged_entities(step_count, ims_instance, mocker):
 
 
 def test_context_manager(ims_instance, mocker):
-    mock_post = mocker.patch('inventory_provider.db.ims.requests.post')
-    mock_post.return_value.status_code = 200
-    mock_post.return_value.text = 'dummy_bt'
-    mocked_get = mocker.patch('inventory_provider.db.ims.requests.get')
+    mocked_get = mocker.patch.object(ims_instance.session, 'get')
     mocked_get.return_value.status_code = 200
     mocked_get.return_value.json.return_value = []
 
@@ -174,17 +171,15 @@ def test_context_manager(ims_instance, mocker):
 
 
 def test_navigation_properties(ims_instance, mocker):
-    mocked_get = mocker.patch('inventory_provider.db.ims.requests.get')
+    mocked_get = mocker.patch.object(ims_instance.session, 'get')
 
     ims_instance.get_entity_by_id('Node', 1234, navigation_properties=[1, 2, 3])
     mocked_get.assert_called_with(
-        f'{BASE_URL}{IMS.IMS_PATH}/Node/1234',
-        headers={'Authorization': 'Bearer dummy_bt'},
-        params={'navigationproperty': 6}, verify=False)
+        f'{BASE_URL}{IMS.IMS_PATH}/Node/1234', params={'navigationproperty': 6})
 
 
 def test_get_entities_gateway_error(ims_instance, mocker):
-    mocked_get = mocker.patch('inventory_provider.db.ims.requests.get')
+    mocked_get = mocker.patch.object(ims_instance.session, 'get')
     mocker.patch('inventory_provider.db.ims.time')
 
     error_message = 'Dummy Gateway Error'