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, 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: 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)