From 6cf4508e2e52da74b5ad4d60f8a8039729e7d596 Mon Sep 17 00:00:00 2001
From: Zach Moody <zmoody@digitalocean.com>
Date: Mon, 5 Apr 2021 23:00:28 -0500
Subject: [PATCH] Fix up DetailEndpoint.list() to continue returning lists

---
 pynetbox/core/endpoint.py      |  6 ++++--
 pynetbox/core/query.py         |  4 ++++
 pynetbox/models/dcim.py        | 16 +++++++---------
 tests/integration/test_dcim.py |  3 ++-
 4 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/pynetbox/core/endpoint.py b/pynetbox/core/endpoint.py
index 281c348d..2031cbdb 100644
--- 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
index c1e126d0..6bd93d44 100644
--- 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
index 90418035..53742969 100644
--- 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
index 91f09756..53fedad0 100644
--- 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):