From 62bee49c8714951a67a81904ad86b171c80eabaf Mon Sep 17 00:00:00 2001 From: Zach Moody Date: Thu, 8 Apr 2021 19:01:29 -0500 Subject: [PATCH 1/3] Implement iterator protocol Make RecordSet a well-behaved iterator. --- pynetbox/core/response.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/pynetbox/core/response.py b/pynetbox/core/response.py index 081b377d..547739e7 100644 --- a/pynetbox/core/response.py +++ b/pynetbox/core/response.py @@ -78,16 +78,16 @@ def __init__(self, endpoint, request, **kwargs): self._response_cache = [] def __iter__(self): + return self + + def __next__(self): if self._response_cache: - yield self.endpoint.return_obj( + return self.endpoint.return_obj( self._response_cache.pop(), self.endpoint.api, self.endpoint ) - for i in self.response: - yield self.endpoint.return_obj(i, self.endpoint.api, self.endpoint) - - def __next__(self): - for i in self: - return i + return self.endpoint.return_obj( + next(self.response), self.endpoint.api, self.endpoint + ) def __len__(self): try: From 118b6b28c7f82cfc6e1979d9be01af8ab96542ca Mon Sep 17 00:00:00 2001 From: Zach Moody Date: Thu, 8 Apr 2021 19:05:38 -0500 Subject: [PATCH 2/3] Re-implement ValueError d06ca91 removed the exeception that was raised when .get() used with kwargs returned more than one item. Fixes #362 --- pynetbox/core/endpoint.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/pynetbox/core/endpoint.py b/pynetbox/core/endpoint.py index ab1359dc..54236d43 100644 --- a/pynetbox/core/endpoint.py +++ b/pynetbox/core/endpoint.py @@ -133,7 +133,19 @@ def get(self, *args, **kwargs): key = None if not key: - return next(self.filter(**kwargs), None) + resp = self.filter(**kwargs) + ret = next(resp, None) + if not ret: + return ret + try: + next(resp) + raise ValueError( + "get() returned more than one result. " + "Check that the kwarg(s) passed are valid for this " + "endpoint or use filter() or all() instead." + ) + except StopIteration: + return ret req = Request( key=key, From 252fa3ea26a27a3016dd406ed92466af7f8065b7 Mon Sep 17 00:00:00 2001 From: Zach Moody Date: Thu, 8 Apr 2021 19:07:27 -0500 Subject: [PATCH 3/3] Add unit tests for Endpoint.get() --- tests/unit/test_endpoint.py | 33 +++++++++++++++++++++++++++++++++ 1 file changed, 33 insertions(+) diff --git a/tests/unit/test_endpoint.py b/tests/unit/test_endpoint.py index 6726352c..bbae08e3 100644 --- a/tests/unit/test_endpoint.py +++ b/tests/unit/test_endpoint.py @@ -59,3 +59,36 @@ def test_choices(self): choices = test_obj.choices() self.assertEqual(choices["letter"][1]["display_name"], "B") self.assertEqual(choices["letter"][1]["value"], 2) + + def test_get_with_filter(self): + with patch( + "pynetbox.core.query.Request._make_call", return_value=Mock() + ) as mock: + mock.return_value = [{"id": 123}] + api = Mock(base_url="http://localhost:8000/api") + app = Mock(name="test") + test_obj = Endpoint(api, app, "test") + test = test_obj.get(name="test") + self.assertEqual(test.id, 123) + + def test_get_greater_than_one(self): + with patch( + "pynetbox.core.query.Request._make_call", return_value=Mock() + ) as mock: + mock.return_value = [{"id": 123}, {"id": 321}] + api = Mock(base_url="http://localhost:8000/api") + app = Mock(name="test") + test_obj = Endpoint(api, app, "test") + with self.assertRaises(ValueError) as _: + test_obj.get(name="test") + + def test_get_no_results(self): + with patch( + "pynetbox.core.query.Request._make_call", return_value=Mock() + ) as mock: + mock.return_value = [] + api = Mock(base_url="http://localhost:8000/api") + app = Mock(name="test") + test_obj = Endpoint(api, app, "test") + test = test_obj.get(name="test") + self.assertIsNone(test)