Merge pull request #352 from digitalocean/fix-detail-list

Fix up DetailEndpoint.list() to continue returning lists
netbox-community · Apr 6, 2021 · dfccaba · dfccaba
2 parents 05fef9d + 6cf4508
commit dfccaba
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 12 deletions.
diff --git a/pynetbox/core/endpoint.py b/pynetbox/core/endpoint.py
@@ -413,10 +413,12 @@ def list(self, **kwargs):
         req = Request(**self.request_kwargs).get(add_params=kwargs)
 
         if self.custom_return:
-            for i in req:
-                yield self.custom_return(
+            return [
+                self.custom_return(
                     i, self.parent_obj.endpoint.api, self.parent_obj.endpoint
                 )
+                for i in req
+            ]
         return req
 
     def create(self, data=None):

diff --git a/pynetbox/core/query.py b/pynetbox/core/query.py
@@ -346,6 +346,10 @@ def get(self, add_params=None):
                 first_run = False
                 for i in req["results"]:
                     yield i
+        elif isinstance(req, list):
+            self.count = len(req)
+            for i in req:
+                yield i
         else:
             self.count = len(req)
             yield req

diff --git a/pynetbox/models/dcim.py b/pynetbox/models/dcim.py
@@ -24,15 +24,13 @@
 
 class TraceableRecord(Record):
     def trace(self):
-        req = list(
-            Request(
-                key=str(self.id) + "/trace",
-                base=self.endpoint.url,
-                token=self.api.token,
-                session_key=self.api.session_key,
-                http_session=self.api.http_session,
-            ).get()
-        )[0]
+        req = Request(
+            key=str(self.id) + "/trace",
+            base=self.endpoint.url,
+            token=self.api.token,
+            session_key=self.api.session_key,
+            http_session=self.api.http_session,
+        ).get()
         uri_to_obj_class_map = {
             "dcim/cables": Cables,
             "dcim/front-ports": FrontPorts,

diff --git a/tests/integration/test_dcim.py b/tests/integration/test_dcim.py
@@ -100,7 +100,8 @@ def init(self, request, rack):
 
     def test_get_elevation(self):
         test = self.fixture.elevation.list()
-        assert next(test)
+        assert test
+        assert isinstance(test, list)
 
 
 class TestManufacturer(BaseTest):