Merge pull request #25 from synkd/add_subscription_allocations_property

Add subscription_allocations property

manifester/helpers.py

@@ -5,6 +5,9 @@
 
 from logzero import logger
 
+MAX_RESULTS_PER_PAGE = 50
+RESULTS_LIMIT = 10000
+
 
 def simple_retry(cmd, cmd_args=None, cmd_kwargs=None, max_timeout=240, _cur_timeout=1):
     """Re(Try) a function given its args and kwargs up until a max timeout."""
@@ -48,6 +51,75 @@ def process_sat_version(sat_version, valid_sat_versions):
     return sat_version
 
 
+def fetch_paginated_data(manifester, endpoint):
+    """Fetch data from the API and account for pagination in the API response.
+
+    Currently used only for subscription allocations and subscription pools.
+    """
+    if endpoint == "allocations":
+        _endpoint_url = manifester.allocations_url
+        _endpoint_data = manifester._allocations
+    elif endpoint == "pools":
+        _endpoint_url = f"{manifester.allocations_url}/{manifester.allocation_uuid}/pools"
+        _endpoint_data = manifester._subscription_pools
+    else:
+        raise ValueError(
+            f"Received value {endpoint} for endpoint argument. Valid values "
+            "for endpoint are 'allocations' or 'pools'."
+        )
+    if not _endpoint_data:
+        _offset = 0
+        data = {
+            "headers": {"Authorization": f"Bearer {manifester.access_token}"},
+            "proxies": manifester.manifest_data.get("proxies"),
+            "params": {
+                "offset": _offset,
+                "limit": RESULTS_LIMIT,
+            },
+        }
+        _endpoint_data = simple_retry(
+            manifester.requester.get,
+            cmd_args=[f"{_endpoint_url}"],
+            cmd_kwargs=data,
+        ).json()
+        if manifester.is_mock and endpoint == "pools":
+            _endpoint_data = _endpoint_data.pool_response
+        elif manifester.is_mock and endpoint == "allocations":
+            _endpoint_data = _endpoint_data.allocations_response
+        _results = len(_endpoint_data["body"])
+        # The endpoints used in the above API call can return a maximum of 50 results. For
+        # organizations with more than 50 subscription allocations or pools, the loop below works
+        # around this limit by repeating calls with a progressively larger value for the `offset`
+        # parameter.
+        while _results == MAX_RESULTS_PER_PAGE:
+            _offset += 50
+            logger.debug(f"Fetching additional data with an offset of {_offset}.")
+            data = {
+                "headers": {"Authorization": f"Bearer {manifester.access_token}"},
+                "proxies": manifester.manifest_data.get("proxies"),
+                "params": {"offset": _offset, "limit": RESULTS_LIMIT},
+            }
+            offset_data = simple_retry(
+                manifester.requester.get,
+                cmd_args=[f"{_endpoint_url}"],
+                cmd_kwargs=data,
+            ).json()
+            if manifester.is_mock and endpoint == "pools":
+                offset_data = offset_data.pool_response
+            elif manifester.is_mock and endpoint == "allocations":
+                offset_data = offset_data.allocations_response
+            _endpoint_data["body"] += offset_data["body"]
+            _results = len(offset_data["body"])
+            total_results = len(_endpoint_data["body"])
+            logger.debug(f"Total {endpoint} available on this account: {total_results}")
+    if endpoint == "allocations":
+        return [
+            a for a in _endpoint_data["body"] if a["name"].startswith(manifester.username_prefix)
+        ]
+    elif endpoint == "pools":
+        return _endpoint_data
+
+
 def fake_http_response_code(good_codes=None, bad_codes=None, fail_rate=0):
     """Return an HTTP response code randomly selected from sets of good and bad codes."""
     if random.random() > (fail_rate / 100):

manifester/manifester.py

@@ -12,9 +12,12 @@
 from logzero import logger
 from requests.exceptions import Timeout
 
-from manifester.helpers import process_sat_version, simple_retry
+from manifester.helpers import fetch_paginated_data, process_sat_version, simple_retry
 from manifester.settings import settings
 
+MAX_RESULTS_PER_PAGE = 50
+RESULTS_LIMIT = 10000
+
 
 class Manifester:
     """Main Manifester class responsible for generating a manifest from the provided settings."""
@@ -32,7 +35,8 @@ def __init__(self, manifest_category, allocation_name=None, **kwargs):
 
             self.requester = requests
             self.is_mock = False
-        self.allocation_name = allocation_name or f"{settings.username_prefix}-" + "".join(
+        self.username_prefix = settings.username_prefix or self.manifest_data.username_prefix
+        self.allocation_name = allocation_name or f"{self.username_prefix}-" + "".join(
             random.sample(string.ascii_letters, 8)
         )
         self.manifest_name = Path(f"{self.allocation_name}_manifest.zip")
@@ -49,6 +53,7 @@ def __init__(self, manifest_category, allocation_name=None, **kwargs):
         self.token_request_url = self.manifest_data.get("url").get("token_request")
         self.allocations_url = self.manifest_data.get("url").get("allocations")
         self._access_token = None
+        self._allocations = None
         self._subscription_pools = None
         self._active_pools = []
         self.sat_version = process_sat_version(
@@ -93,6 +98,19 @@ def valid_sat_versions(self):
         valid_sat_versions = [ver_dict["value"] for ver_dict in sat_versions_response["body"]]
         return valid_sat_versions
 
+    @property
+    def subscription_allocations(self):
+        """Representation of subscription allocations in an account.
+
+        Filtered by username_prefix.
+        """
+        return fetch_paginated_data(self, "allocations")
+
+    @property
+    def subscription_pools(self):
+        """Reprentation of subscription pools in an account."""
+        return fetch_paginated_data(self, "pools")
+
     def create_subscription_allocation(self):
         """Creates a new consumer in the provided RHSM account and returns its UUID."""
         allocation_data = {
@@ -144,55 +162,6 @@ def delete_subscription_allocation(self):
         )
         return response
 
-    @property
-    def subscription_pools(self):
-        """Fetches the list of subscription pools from account.
-
-        Returns a list of dictionaries containing metadata from the pools.
-        """
-        MAX_RESULTS_PER_PAGE = 50
-        if not self._subscription_pools:
-            _offset = 0
-            data = {
-                "headers": {"Authorization": f"Bearer {self.access_token}"},
-                "proxies": self.manifest_data.get("proxies"),
-                "params": {"offset": _offset},
-            }
-            self._subscription_pools = simple_retry(
-                self.requester.get,
-                cmd_args=[f"{self.allocations_url}/{self.allocation_uuid}/pools"],
-                cmd_kwargs=data,
-            ).json()
-            if self.is_mock:
-                self._subscription_pools = self._subscription_pools.pool_response
-            _results = len(self._subscription_pools["body"])
-            # The endpoint used in the above API call can return a maximum of 50 subscription pools.
-            # For organizations with more than 50 subscription pools, the loop below works around
-            # this limit by repeating calls with a progressively larger value for the `offset`
-            # parameter.
-            while _results == MAX_RESULTS_PER_PAGE:
-                _offset += 50
-                logger.debug(f"Fetching additional subscription pools with an offset of {_offset}.")
-                data = {
-                    "headers": {"Authorization": f"Bearer {self.access_token}"},
-                    "proxies": self.manifest_data.get("proxies"),
-                    "params": {"offset": _offset},
-                }
-                offset_pools = simple_retry(
-                    self.requester.get,
-                    cmd_args=[f"{self.allocations_url}/{self.allocation_uuid}/pools"],
-                    cmd_kwargs=data,
-                ).json()
-                if self.is_mock:
-                    offset_pools = offset_pools.pool_response
-                self._subscription_pools["body"] += offset_pools["body"]
-                _results = len(offset_pools["body"])
-                total_pools = len(self._subscription_pools["body"])
-                logger.debug(
-                    f"Total subscription pools available for this allocation: {total_pools}"
-                )
-        return self._subscription_pools
-
     def add_entitlements_to_allocation(self, pool_id, entitlement_quantity):
         """Attempts to add the set of subscriptions defined in the settings to the allocation."""
         data = {

pyproject.toml

@@ -179,5 +179,8 @@ mark-parentheses = false
 [tool.ruff.flake8-quotes]
 inline-quotes = "single"
 
+[tool.ruff.lint.pylint]
+max-branches = 16
+
 [tool.ruff.mccabe]
 max-complexity = 20

tests/test_manifester.py

@@ -13,6 +13,7 @@
     "log_level": "debug",
     "offline_token": "test",
     "proxies": {"https": ""},
+    "username_prefix": "test_user",
     "url": {
         "token_request": "https://sso.redhat.com/auth/realms/redhat-external/protocol/openid-connect/token",
         "allocations": "https://api.access.redhat.com/management/v1/allocations",
@@ -64,6 +65,15 @@
     ],
 }
 
+sub_allocations_response = {
+    "body": [
+        {
+            "uuid": f"{uuid.uuid4().hex}",
+            "name": "test_user-" + "".join(random.sample(string.ascii_letters, 8)),
+        }
+    ]
+}
+
 
 class RhsmApiStub(MockStub):
     """Returns mock responses for RHSM API endpoints related to creating manifests."""
@@ -122,7 +132,25 @@ def get(self, *args, **kwargs):
             else:
                 self.pool_response["body"] += sub_pool_response["body"]
                 return self
-        if "allocations" in args[0] and not ("export" in args[0] or "pools" in args[0]):
+        if args[0].endswith("allocations") and self._has_offset:
+            if kwargs["params"]["offset"] != 50:
+                self.allocations_response = {"body": []}
+                for _x in range(50):
+                    self.allocations_response["body"].append(
+                        {
+                            "uuid": f"{uuid.uuid4().hex}",
+                            "name": f'{"".join(random.sample(string.ascii_letters, 12))}',
+                        }
+                    )
+                return self
+            else:
+                self.allocations_response["body"] += sub_allocations_response["body"]
+                return self
+        if (
+            "allocations" in args[0]
+            and not ("export" in args[0] or "pools" in args[0])
+            and not self._has_offset
+        ):
             self.allocation_data = "this allocation data also includes entitlement data"
             return self
         if args[0].endswith("export"):
@@ -226,6 +254,15 @@ def test_get_subscription_pools_with_offset():
     assert len(manifester.subscription_pools["body"]) > 50
 
 
+def test_subscription_allocation_username_prefix_filter():
+    """Test that all allocations in subscription_allocations property matching username prefix."""
+    manifester = Manifester(
+        manifest_category=manifest_data, requester=RhsmApiStub(in_dict=None, has_offset=True)
+    )
+    for allocation in manifester.subscription_allocations:
+        assert allocation["name"].startswith(manifest_data["username_prefix"])
+
+
 def test_correct_subs_added_to_allocation():
     """Test that subs added to the allocation match the subscription data in manifester's config."""
     manifester = Manifester(manifest_category=manifest_data, requester=RhsmApiStub(in_dict=None))