Minor adjustments

Author: samwelkanda

msgraphcore/graph_client.py

@@ -2,9 +2,8 @@
 from msgraphcore.middleware.request_context import RequestContext
 from typing import List, Optional
 
-from requests import Session
-from requests import PreparedRequest
-from requests import Request
+from requests import Request, Session
+
 
 from msgraphcore.client_factory import HTTPClientFactory
 from msgraphcore.middleware.abc_token_credential import TokenCredential
@@ -15,16 +14,15 @@
 
 def attach_context(func):
     def wrapper(*args, **kwargs):
-        mw_control = dict()
-        headers = kwargs.pop('headers', {})
-
+        middleware_control = dict()
+        
         for option in supported_options:
-            op = kwargs.pop(option, None)
-
-            if op:
-                mw_control.update({option: op})
-
-        req_context = RequestContext(mw_control, headers)
+            value = kwargs.pop(option, None)
+            if value:
+                middleware_control.update({option: value})
+                
+        headers = kwargs.get('headers', {})
+        req_context = RequestContext(middleware_control, headers)
 
         request = func(*args, **kwargs)
         request.context = req_context
@@ -79,7 +77,8 @@ def put(self, url, data=None, **kwargs):
         :param \*\*kwargs: Optional arguments that ``request`` takes.
         :rtype: requests.Response
         """
-        return self.graph_session.put(self._graph_url(url), data, **kwargs)
+        prepped_req = self.prepare_request('PUT', self._graph_url(url), **kwargs)
+        return self.graph_session.send(prepped_req)
 
     def patch(self, url, data=None, **kwargs):
         r"""Sends a PATCH request. Returns :class:`Response` object.
@@ -89,15 +88,17 @@ def patch(self, url, data=None, **kwargs):
         :param \*\*kwargs: Optional arguments that ``request`` takes.
         :rtype: requests.Response
         """
-        return self.graph_session.patch(self._graph_url(url), data, **kwargs)
+        prepped_req = self.prepare_request('PATCH', self._graph_url(url), **kwargs)
+        return self.graph_session.send(prepped_req)
 
     def delete(self, url, **kwargs):
         r"""Sends a DELETE request. Returns :class:`Response` object.
         :param url: URL for the new :class:`Request` object.
         :param \*\*kwargs: Optional arguments that ``request`` takes.
         :rtype: requests.Response
         """
-        return self.graph_session.delete(self._graph_url(url), **kwargs)
+        prepped_req = self.prepare_request('DELETE', self._graph_url(url), **kwargs)
+        return self.graph_session.send(prepped_req)
 
     def _graph_url(self, url: str) -> str:
         """Appends BASE_URL to user provided path

msgraphcore/middleware/middleware.py

@@ -25,6 +25,11 @@ def add_middleware(self, middleware):
             self._middleware = middleware
 
     def send(self, request, **kwargs):
+        
+        if not hasattr(request, 'context'):
+            headers = request.headers
+            request.context = RequestContext(dict(), headers)
+        
         if self._middleware_present():
             return self._middleware.send(request, **kwargs)
         # No middleware in pipeline, call superclass' send

msgraphcore/middleware/request_context.py

@@ -16,5 +16,4 @@ def feature_usage(self):
 
     @feature_usage.setter
     def set_feature_usage(self, flag: FeatureUsageFlag):
-        print(type(self._feature_usage))
         self._feature_usage = self._feature_usage | flag