From b13aee3257ecb38231629bac323f908e2f1dca36 Mon Sep 17 00:00:00 2001
From: Vladimir Sedmik <vsedmik@redhat.com>
Date: Wed, 18 Oct 2023 14:12:58 +0200
Subject: [PATCH] Add support for capsule content counts

---
 nailgun/entities.py    | 45 +++++++++++++++++++++++++++++++++++++++++-
 tests/test_entities.py |  9 ++++++++-
 2 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/nailgun/entities.py b/nailgun/entities.py
index d9f2caa7..8a7eb682 100644
--- a/nailgun/entities.py
+++ b/nailgun/entities.py
@@ -837,6 +837,10 @@ def __init__(self, server_config=None, **kwargs):
             ),
             'organization': entity_fields.OneToManyField(Organization),
             'url': entity_fields.StringField(required=True),
+            'hosts_count': entity_fields.IntegerField(),
+            'download_policy': entity_fields.StringField(),
+            'supported_pulp_types': entity_fields.StringField(),
+            'lifecycle_environments': entity_fields.StringField(),
         }
         self._meta = {
             'api_path': 'katello/api/capsules',
@@ -943,6 +947,42 @@ def content_get_sync(self, synchronous=True, timeout=None, **kwargs):
         response = client.get(self.path('content_sync'), **kwargs)
         return _handle_response(response, self._server_config, synchronous, timeout)
 
+    def content_counts(self, synchronous=True, timeout=None, **kwargs):
+        """List content counts for the capsule.
+
+        :param synchronous: What should happen if the server returns an HTTP
+            202 (accepted) status code? Wait for the task to complete if
+            ``True``. Immediately return the server's response otherwise.
+        :param timeout: Maximum number of seconds to wait until timing out.
+            Defaults to ``nailgun.entity_mixins.TASK_TIMEOUT``.
+        :param kwargs: Arguments to pass to requests.
+        :returns: The server's response, with all JSON decoded.
+        :raises: ``requests.exceptions.HTTPError`` If the server responds with
+            an HTTP 4XX or 5XX message.
+        """
+        kwargs = kwargs.copy()
+        kwargs.update(self._server_config.get_client_kwargs())
+        response = client.get(self.path('content_counts'), **kwargs)
+        return _handle_response(response, self._server_config, synchronous, timeout)
+
+    def content_update_counts(self, synchronous=True, timeout=None, **kwargs):
+        """Update content counts for the capsule.
+
+        :param synchronous: What should happen if the server returns an HTTP
+            202 (accepted) status code? Wait for the task to complete if
+            ``True``. Immediately return the server's response otherwise.
+        :param timeout: Maximum number of seconds to wait until timing out.
+            Defaults to ``nailgun.entity_mixins.TASK_TIMEOUT``.
+        :param kwargs: Arguments to pass to requests.
+        :returns: The server's response, with all JSON decoded.
+        :raises: ``requests.exceptions.HTTPError`` If the server responds with
+            an HTTP 4XX or 5XX message.
+        """
+        kwargs = kwargs.copy()
+        kwargs.update(self._server_config.get_client_kwargs())
+        response = client.post(self.path('content_update_counts'), **kwargs)
+        return _handle_response(response, self._server_config, synchronous, timeout)
+
     def path(self, which=None):
         """Extend ``nailgun.entity_mixins.Entity.path``.
 
@@ -952,7 +992,10 @@ def path(self, which=None):
             /capsules/<id>/content/lifecycle_environments
         content_sync
             /capsules/<id>/content/sync
-
+        content_counts
+            /capsules/<id>/content/counts
+        content_update_counts
+            /capsules/<id>/content/update_counts
 
         ``super`` is called otherwise.
 
diff --git a/tests/test_entities.py b/tests/test_entities.py
index 9618d5cd..193d2945 100644
--- a/tests/test_entities.py
+++ b/tests/test_entities.py
@@ -564,7 +564,12 @@ def test_capsule(self):
         * ``Capsule().path('content_sync')``
         """
         capsule = entities.Capsule(self.cfg, id=gen_integer(1, 100))
-        for which in ('content_lifecycle_environments', 'content_sync'):
+        for which in (
+            'content_lifecycle_environments',
+            'content_sync',
+            'content_counts',
+            'content_update_counts',
+        ):
             with self.subTest(which):
                 path = capsule.path(which)
                 which_parts = which.split("_", 1)
@@ -2102,6 +2107,8 @@ def setUpClass(cls):
             (entities.Capsule(**generic).content_get_sync, 'get'),
             (entities.Capsule(**generic).content_lifecycle_environments, 'get'),
             (entities.Capsule(**generic).content_sync, 'post'),
+            (entities.Capsule(**generic).content_counts, 'get'),
+            (entities.Capsule(**generic).content_update_counts, 'post'),
             (entities.Role(**generic).clone, 'post'),
             (entities.ProvisioningTemplate(**generic).build_pxe_default, 'post'),
             (entities.ProvisioningTemplate(**generic).clone, 'post'),