Make custom work with get_entity

Author: Sylvain Fankhauser

pyodata/v2/service.py

@@ -242,6 +242,7 @@ def __init__(self, url, connection, handler, headers=None):
         self._handler = handler
         self._headers = headers or dict()
         self._logger = logging.getLogger(LOGGER_NAME)
+        self._customs = {}  # string -> string hash
 
     @property
     def handler(self):
@@ -256,7 +257,7 @@ def get_path(self):
     def get_query_params(self):
         """Get query params"""
         # pylint: disable=no-self-use
-        return {}
+        return dict(self._customs)
 
     def get_method(self):
         """Get HTTP method"""
@@ -331,6 +332,12 @@ def execute(self):
 
         return self._handler(response)
 
+    def custom(self, name, value):
+        """Adds a custom name-value pair."""
+        # returns QueryRequest
+        self._customs[name] = value
+        return self
+
 
 class EntityGetRequest(ODataHttpRequest):
     """Used for GET operations of a single entity"""
@@ -603,15 +610,8 @@ def __init__(self, url, connection, handler, last_segment):
         self._select = None
         self._expand = None
         self._last_segment = last_segment
-        self._customs = {}  # string -> string hash
         self._logger.debug('New instance of QueryRequest for last segment: %s', self._last_segment)
 
-    def custom(self, name, value):
-        """Adds a custom name-value pair."""
-        # returns QueryRequest
-        self._customs[name] = value
-        return self
-
     def count(self):
         """Sets a flag to return the number of items."""
         self._count = True
@@ -685,9 +685,6 @@ def get_query_params(self):
         if self._select is not None:
             qparams['$select'] = self._select
 
-        for key, val in self._customs.items():
-            qparams[key] = val
-
         if self._expand is not None:
             qparams['$expand'] = self._expand
 

tests/test_service_v2.py

@@ -2180,3 +2180,18 @@ def test_odata_http_response():
     assert isinstance(response.headers, dict)
     assert response.headers['Content-Type'] == 'application/json'
     assert response.json()['d']['ID'] == 23
+
+
+@responses.activate
+def test_custom_with_get_entity(service):
+    """ Test that `custom` can be called after `get_entity`"""
+
+    responses.add(
+        responses.GET,
+        "{0}/MasterEntities('12345')?foo=bar".format(service.url),
+        headers={'Content-type': 'application/json'},
+        json={'d': {'Key': '12345'}},
+        status=200)
+
+    entity = service.entity_sets.MasterEntities.get_entity('12345').custom("foo", "bar").execute()
+    assert entity.Key == '12345'