Merge pull request #253 from digitalocean/record-return-custom

Record return custom

Author: Zach Moody

pynetbox/core/endpoint.py

@@ -19,6 +19,15 @@
 RESERVED_KWARGS = ("id", "pk", "limit", "offset")
 
 
+def response_loader(req, return_obj, endpoint):
+    if isinstance(req, list):
+        return [
+            return_obj(i, endpoint.api, endpoint)
+            for i in req
+        ]
+    return return_obj(req, endpoint.api, endpoint)
+
+
 class Endpoint(object):
     """Represent actions available on endpoints in the Netbox API.
 
@@ -73,9 +82,6 @@ def _lookup_ret_obj(self, name, model):
             ret = Record
         return ret
 
-    def _response_loader(self, values):
-        return self.return_obj(values, self.api, self)
-
     def all(self):
         """Queries the 'ListView' of a given endpoint.
 
@@ -97,7 +103,7 @@ def all(self):
             threading=self.api.threading,
         )
 
-        return [self._response_loader(i) for i in req.get()]
+        return response_loader(req.get(), self.return_obj, self)
 
     def get(self, *args, **kwargs):
         r"""Queries the DetailsView of a given endpoint.
@@ -157,7 +163,7 @@ def get(self, *args, **kwargs):
         except RequestError:
             return None
 
-        return self._response_loader(req.get())
+        return response_loader(req.get(), self.return_obj, self)
 
     def filter(self, *args, **kwargs):
         r"""Queries the 'ListView' of a given endpoint.
@@ -223,8 +229,7 @@ def filter(self, *args, **kwargs):
             threading=self.api.threading,
         )
 
-        ret = [self._response_loader(i) for i in req.get()]
-        return ret
+        return response_loader(req.get(), self.return_obj, self)
 
     def create(self, *args, **kwargs):
         r"""Creates an object on an endpoint.
@@ -284,10 +289,7 @@ def create(self, *args, **kwargs):
             http_session=self.api.http_session,
         ).post(args[0] if args else kwargs)
 
-        if isinstance(req, list):
-            return [self._response_loader(i) for i in req]
-
-        return self._response_loader(req)
+        return response_loader(req, self.return_obj, self)
 
     def choices(self):
         """ Returns all choices from the endpoint.
@@ -429,15 +431,8 @@ def list(self, **kwargs):
         req = Request(**self.request_kwargs).get(add_params=kwargs)
 
         if self.custom_return:
-            if isinstance(req, list):
-                return [
-                    self.custom_return(
-                        i, self.parent_obj.api, self.parent_obj.endpoint
-                    )
-                    for i in req
-                ]
-            return self.custom_return(
-                req, self.parent_obj.api, self.parent_obj.endpoint
+            return response_loader(
+                req, self.custom_return, self.parent_obj.endpoint
             )
         return req
 
@@ -454,9 +449,13 @@ def create(self, data=None):
         :returns: A dictionary or list of dictionaries its created in
             NetBox.
         """
-        if not data:
-            return Request(**self.request_kwargs).post({})
-        return Request(**self.request_kwargs).post(data)
+        data = data or {}
+        req = Request(**self.request_kwargs).post(data)
+        if self.custom_return:
+            return response_loader(
+                req, self.custom_return, self.parent_obj.endpoint
+            )
+        return req
 
 
 class RODetailEndpoint(DetailEndpoint):

pynetbox/models/ipam.py

@@ -55,7 +55,7 @@ def available_ips(self):
         >>> len(create)
         2
         """
-        return DetailEndpoint(self, "available-ips")
+        return DetailEndpoint(self, "available-ips", custom_return=IpAddresses)
 
     @property
     def available_prefixes(self):
@@ -86,7 +86,9 @@ def available_prefixes(self):
         u'10.1.1.56/29'
 
         """
-        return DetailEndpoint(self, "available-prefixes")
+        return DetailEndpoint(
+            self, "available-prefixes", custom_return=Prefixes
+        )
 
 
 class Aggregates(Record):

tests/test_ipam.py

@@ -162,18 +162,6 @@ def test_get_available_ips(self, mock):
         return_value=Response(fixture='ipam/prefix.json'),
     )
     def test_create_available_ips(self, _, post):
-        expected_result = {
-            'status': 1,
-            'description': '',
-            'nat_inside': None,
-            'role': None,
-            'vrf': None,
-            'address':
-            '10.1.1.1/32',
-            'interface': None,
-            'id': 1,
-            'tenant': None
-        }
         create_parms = dict(
             status=2,
         )
@@ -186,7 +174,7 @@ def test_create_available_ips(self, _, post):
             json=create_parms,
         )
         self.assertTrue(ret)
-        self.assertEqual(ret, expected_result)
+        self.assertTrue(isinstance(ret, pynetbox.models.ipam.IpAddresses))
 
     @patch(
         'pynetbox.core.query.requests.sessions.Session.get',
@@ -227,6 +215,7 @@ def test_create_available_prefixes(self, _, post):
             json=create_parms,
         )
         self.assertTrue(ret)
+        self.assertTrue(isinstance(ret[0], pynetbox.models.ipam.Prefixes))
 
 
 class IPAddressTestCase(Generic.Tests):