Merge pull request #112 from digitalocean/106-ru-detail-endpoint

Fixes #106 - Rack Unit DetailEndpoint

Author: Zach Moody

pynetbox/dcim.py

@@ -60,10 +60,7 @@ def napalm(self):
         {"get_facts": {"interface_list": ["ge-0/0/0"]}}
 
         """
-        return RODetailEndpoint(
-            'napalm',
-            parent_obj=self,
-        )
+        return RODetailEndpoint(self, 'napalm')
 
 
 class InterfaceConnections(Record):
@@ -92,3 +89,29 @@ class VirtualChassis(Record):
 
     def __str__(self):
         return self.master.display_name
+
+
+class RUs(Record):
+
+    device = Devices
+
+
+class Racks(Record):
+
+    @property
+    def units(self):
+        """ Represents the ``units`` detail endpoint.
+
+        Returns a DetailEndpoint object that is the interface for
+        viewing response from the units endpoint.
+
+        :returns: :py:class:`.DetailEndpoint`
+
+        :Examples:
+
+        >>> rack = nb.dcim.racks.get(123)
+        >>> rack.units.list()
+        {"get_facts": {"interface_list": ["ge-0/0/0"]}}
+
+        """
+        return RODetailEndpoint(self, 'units', custom_return=RUs)

pynetbox/ipam.py

@@ -57,10 +57,7 @@ def available_ips(self):
         >>> len(create)
         2
         """
-        return DetailEndpoint(
-            'available-ips',
-            parent_obj=self,
-        )
+        return DetailEndpoint(self, 'available-ips')
 
     @property
     def available_prefixes(self):
@@ -91,10 +88,7 @@ def available_prefixes(self):
         u'10.1.1.56/29'
 
         '''
-        return DetailEndpoint(
-            'available-prefixes',
-            parent_obj=self,
-        )
+        return DetailEndpoint(self, 'available-prefixes')
 
 
 class Aggregates(IPRecord):

pynetbox/lib/endpoint.py

@@ -305,22 +305,20 @@ class DetailEndpoint(object):
     traditional endpoints are handled with this class.
     '''
 
-    def __init__(self, name, parent_obj=None):
-        self.token = parent_obj.api_kwargs.get('token')
-        self.version = parent_obj.api_kwargs.get('version')
-        self.session_key = parent_obj.api_kwargs.get('session_key')
-        self.ssl_verify = parent_obj.api_kwargs.get('ssl_verify')
+    def __init__(self, parent_obj, name, custom_return=None):
+        self.parent_obj = parent_obj
+        self.custom_return = custom_return
         self.url = "{}/{}/{}/".format(
             parent_obj.endpoint_meta.get('url'),
             parent_obj.id,
             name
         )
         self.request_kwargs = dict(
             base=self.url,
-            token=self.token,
-            session_key=self.session_key,
-            version=self.version,
-            ssl_verify=self.ssl_verify,
+            token=parent_obj.api_kwargs.get('token'),
+            session_key=parent_obj.api_kwargs.get('session_key'),
+            version=parent_obj.api_kwargs.get('version'),
+            ssl_verify=parent_obj.api_kwargs.get('ssl_verify'),
         )
 
     def list(self, **kwargs):
@@ -333,16 +331,32 @@ def list(self, **kwargs):
             E.g. ``.list(method='get_facts')`` would be converted to
             ``.../?method=get_facts``.
 
-        :returns: A dictionary or list of dictionaries its retrieved
-            from NetBox.
+        :returns: A dictionary or list of dictionaries retrieved from
+            NetBox.
         """
         if kwargs:
             self.request_kwargs['base'] = '{}{}'.format(
                 self.url,
                 url_param_builder(kwargs)
             )
-
-        return Request(**self.request_kwargs).get()
+        req = Request(**self.request_kwargs).get()
+
+        if self.custom_return:
+            if isinstance(req, list):
+                return [
+                    self.custom_return(
+                        i,
+                        api_kwargs=self.parent_obj.api_kwargs,
+                        endpoint_meta=self.parent_obj.endpoint_meta
+                    )
+                    for i in req
+                ]
+            return self.custom_return(
+                req,
+                api_kwargs=self.parent_obj.api_kwargs,
+                endpoint_meta=self.parent_obj.endpoint_meta
+            )
+        return req
 
     def create(self, data={}):
         """The write operation for a detail endpoint.

tests/fixtures/dcim/rack_u.json

@@ -0,0 +1,35 @@
+{
+    "count": 3,
+    "next": null,
+    "previous": null,
+    "results": [
+        {
+            "id": 48,
+            "name": "U1",
+            "face": 0,
+            "device": {
+                "id": 130,
+                "url": "http://localhost:8000/api/dcim/devices/1/",
+                "name": "tst-device1",
+                "display_name": "tst-device1"
+            }
+        },
+        {
+            "id": 47,
+            "name": "U2",
+            "face": 0,
+            "device": null
+        },
+        {
+            "id": 46,
+            "name": "U3",
+            "face": 0,
+            "device": {
+                "id": 1859,
+                "url": "http://localhost:8000/api/dcim/devices/2/",
+                "name": "tst-device2",
+                "display_name": "tst-device2"
+            }
+        }
+    ]
+}

tests/test_dcim.py

@@ -381,6 +381,26 @@ def test_get_all(self, mock):
 class RackTestCase(unittest.TestCase, GenericTest):
     name = 'racks'
 
+    @patch(
+        'pynetbox.lib.query.requests.get',
+        side_effect=[
+            Response(fixture='dcim/rack.json'),
+            Response(fixture='dcim/rack_u.json'),
+        ]
+    )
+    def test_get_units(self, mock):
+        test = nb.racks.get(1)
+        ret = test.units.list()
+        mock.assert_called_with(
+            'http://localhost:8000/api/dcim/racks/1/units/',
+            headers=HEADERS,
+            verify=True,
+        )
+        self.assertTrue(ret)
+        self.assertTrue(
+            isinstance(ret[0].device, pynetbox.dcim.Devices)
+        )
+
 
 class RackRoleTestCase(unittest.TestCase, GenericTest):
     name = 'rack_roles'