From 74cba023637bef3131ad4a5f97494bb1a7475095 Mon Sep 17 00:00:00 2001 From: Calvin Remsburg Date: Thu, 12 Dec 2024 17:01:45 -0600 Subject: [PATCH] Support for auto-pagination of retrievals with list method --- scm/config/objects/address.py | 88 ++++++++++++++---------- tests/factories.py | 87 +++++++++++++++++++---- tests/scm/config/objects/test_address.py | 68 +++++++++++++++++- 3 files changed, 191 insertions(+), 52 deletions(-) diff --git a/scm/config/objects/address.py b/scm/config/objects/address.py index d3695382..fd2ae315 100644 --- a/scm/config/objects/address.py +++ b/scm/config/objects/address.py @@ -23,7 +23,9 @@ class Address(BaseObject): """ ENDPOINT = "/config/objects/v1/addresses" - DEFAULT_LIMIT = 10000 + MAX_LIMIT = ( + 5000 # The maximum number of objects returned by the API in a single request + ) def __init__( self, @@ -237,8 +239,6 @@ def list( }, ) - params = {"limit": self.DEFAULT_LIMIT} - container_parameters = self._build_container_params( folder, snippet, @@ -253,48 +253,64 @@ def list( details={"error": "Invalid container parameters"}, ) - params.update(container_parameters) + # Pagination logic + limit = self.MAX_LIMIT + offset = 0 + all_addresses = [] - response = self.api_client.get( - self.ENDPOINT, - params=params, - ) + while True: + params = container_parameters.copy() + params["limit"] = limit + params["offset"] = offset - if not isinstance(response, dict): - raise InvalidObjectError( - message="Invalid response format: expected dictionary", - error_code="E003", - http_status_code=500, - details={"error": "Response is not a dictionary"}, + response = self.api_client.get( + self.ENDPOINT, + params=params, ) - if "data" not in response: - raise InvalidObjectError( - message="Invalid response format: missing 'data' field", - error_code="E003", - http_status_code=500, - details={ - "field": "data", - "error": '"data" field missing in the response', - }, - ) + if not isinstance(response, dict): + raise InvalidObjectError( + message="Invalid response format: expected dictionary", + error_code="E003", + http_status_code=500, + details={"error": "Response is not a dictionary"}, + ) - if not isinstance(response["data"], list): - raise InvalidObjectError( - message="Invalid response format: 'data' field must be a list", - error_code="E003", - http_status_code=500, - details={ - "field": "data", - "error": '"data" field must be a list', - }, - ) + if "data" not in response: + raise InvalidObjectError( + message="Invalid response format: missing 'data' field", + error_code="E003", + http_status_code=500, + details={ + "field": "data", + "error": '"data" field missing in the response', + }, + ) + + if not isinstance(response["data"], list): + raise InvalidObjectError( + message="Invalid response format: 'data' field must be a list", + error_code="E003", + http_status_code=500, + details={ + "field": "data", + "error": '"data" field must be a list', + }, + ) + + data = response["data"] + addresses_chunk = [AddressResponseModel(**item) for item in data] + all_addresses.extend(addresses_chunk) + + # If we got fewer than 'limit' objects, we've reached the end + if len(data) < limit: + break - addresses = [AddressResponseModel(**item) for item in response["data"]] + offset += limit # Apply existing filters first addresses = self._apply_filters( - addresses, + all_addresses, filters, ) diff --git a/tests/factories.py b/tests/factories.py index a5e460b5..91d9291c 100644 --- a/tests/factories.py +++ b/tests/factories.py @@ -231,38 +231,97 @@ class Meta: tag = ["response-tag"] folder = "Texas" - # Address types default to None + # No defaults for ip_netmask, ip_range, ip_wildcard, fqdn here ip_netmask = None ip_range = None ip_wildcard = None fqdn = None @classmethod - def with_ip_netmask(cls, ip_netmask="192.168.1.1/32", **kwargs): - return cls(ip_netmask=ip_netmask, **kwargs) + def with_ip_netmask( + cls, + ip_netmask="192.168.1.1/32", + **kwargs, + ): + # Clears out other fields to ensure only one type is set + return cls( + ip_netmask=ip_netmask, + ip_range=None, + ip_wildcard=None, + fqdn=None, + **kwargs, + ) @classmethod - def with_fqdn(cls, fqdn="example.com", **kwargs): - return cls(fqdn=fqdn, **kwargs) + def with_fqdn( + cls, + fqdn="example.com", + **kwargs, + ): + return cls( + fqdn=fqdn, + ip_netmask=None, + ip_range=None, + ip_wildcard=None, + **kwargs, + ) @classmethod - def with_ip_range(cls, ip_range="192.168.0.1-192.168.0.10", **kwargs): - return cls(ip_range=ip_range, **kwargs) + def with_ip_range( + cls, + ip_range="192.168.0.1-192.168.0.10", + **kwargs, + ): + return cls( + ip_range=ip_range, + ip_netmask=None, + ip_wildcard=None, + fqdn=None, + **kwargs, + ) @classmethod - def with_ip_wildcard(cls, ip_wildcard="10.20.1.0/0.0.248.255", **kwargs): - return cls(ip_wildcard=ip_wildcard, **kwargs) + def with_ip_wildcard( + cls, + ip_wildcard="10.20.1.0/0.0.248.255", + **kwargs, + ): + return cls( + ip_wildcard=ip_wildcard, + ip_netmask=None, + ip_range=None, + fqdn=None, + **kwargs, + ) @classmethod - def with_snippet(cls, **kwargs): - return cls(folder=None, snippet="TestSnippet", **kwargs) + def with_snippet( + cls, + **kwargs, + ): + return cls( + folder=None, + snippet="TestSnippet", + **kwargs, + ) @classmethod - def with_device(cls, **kwargs): - return cls(folder=None, device="TestDevice", **kwargs) + def with_device( + cls, + **kwargs, + ): + return cls( + folder=None, + device="TestDevice", + **kwargs, + ) @classmethod - def from_request(cls, request_model: AddressCreateModel, **kwargs): + def from_request( + cls, + request_model: AddressCreateModel, + **kwargs, + ): """Create a response model based on a request model.""" data = request_model.model_dump() data["id"] = str(uuid.uuid4()) diff --git a/tests/scm/config/objects/test_address.py b/tests/scm/config/objects/test_address.py index 49878081..8a2e0577 100644 --- a/tests/scm/config/objects/test_address.py +++ b/tests/scm/config/objects/test_address.py @@ -78,8 +78,9 @@ def test_list_valid(self): self.mock_scm.get.assert_called_once_with( # noqa "/config/objects/v1/addresses", params={ - "limit": 10000, "folder": "All", + "limit": 5000, + "offset": 0, }, ) assert isinstance(existing_objects, list) @@ -164,8 +165,9 @@ def test_list_filters_valid(self): self.mock_scm.get.assert_called_once_with( # noqa "/config/objects/v1/addresses", params={ - "limit": 10000, "folder": "Texas", + "limit": 5000, + "offset": 0, }, ) @@ -483,6 +485,68 @@ def test_list_server_error(self): assert error_response["_errors"][0]["message"] == "An internal error occurred" assert error_response["_errors"][0]["details"]["errorType"] == "Internal Error" + def test_list_pagination_multiple_pages(self): + """ + Test that the list method correctly aggregates data from multiple pages when + more than 5000 objects are returned. + + We will create a scenario with 12,345 objects: + - First API call returns 5000 objects. + - Second API call returns another 5000 objects. + - Third API call returns the remaining 2345 objects. + """ + total_objects = 12345 + first_batch_count = 5000 + second_batch_count = 5000 + third_batch_count = ( + total_objects - first_batch_count - second_batch_count + ) # 2345 + + # Create batches using the factory in a similar manner to existing tests + first_batch = [ + AddressResponseFactory.with_ip_netmask(ip_netmask=f"10.0.0.{i % 254}/32") + for i in range(first_batch_count) + ] + second_batch = [ + AddressResponseFactory.with_ip_netmask(ip_netmask=f"10.0.1.{i % 254}/32") + for i in range(second_batch_count) + ] + third_batch = [ + AddressResponseFactory.with_ip_netmask(ip_netmask=f"10.0.2.{i % 254}/32") + for i in range(third_batch_count) + ] + # Convert them to dicts for the mock responses + first_page = {"data": [obj.model_dump() for obj in first_batch]} + second_page = {"data": [obj.model_dump() for obj in second_batch]} + third_page = {"data": [obj.model_dump() for obj in third_batch]} + + def mock_get(endpoint, params): + offset = params.get("offset", 0) + if offset == 0: + return first_page + elif offset == 5000: + return second_page + elif offset == 10000: + return third_page + else: + # No more data + return {"data": []} + + # Set our mock to the new side_effect + self.mock_scm.get.side_effect = mock_get + + # Invoke the list method; it should loop through all pages + results = self.client.list(folder="Texas") + + # Validate that we received all objects + assert len(results) == total_objects + assert all(isinstance(r, AddressResponseModel) for r in results) + + # Optional: check first and last objects to ensure correct ordering/consistency + # Since factories usually increment sequences, verify at least that the fields are present + assert hasattr(results[0], "name") + assert hasattr(results[-1], "name") + # -------------------- New Tests for exact_match and Exclusions -------------------- def test_list_exact_match(self):