Merge pull request #232 from opsmill/wvd-20240110-add-filters-for-count-method

Add possibility to use filters in the SDK client's count method

Author: wvandeun

changelog/count-method-filters.changed.md

@@ -0,0 +1 @@
+Added possibility to use filters for the SDK client's count method

infrahub_sdk/client.py

@@ -547,16 +547,21 @@ async def count(
         at: Timestamp | None = None,
         branch: str | None = None,
         timeout: int | None = None,
+        **kwargs: Any,
     ) -> int:
         """Return the number of nodes of a given kind."""
+        filters = kwargs
         schema = await self.schema.get(kind=kind, branch=branch)
 
         branch = branch or self.default_branch
         if at:
             at = Timestamp(at)
 
         response = await self.execute_graphql(
-            query=Query(query={schema.kind: {"count": None}}).render(), branch_name=branch, at=at, timeout=timeout
+            query=Query(query={schema.kind: {"count": None, "@filters": filters}}).render(),
+            branch_name=branch,
+            at=at,
+            timeout=timeout,
         )
         return int(response.get(schema.kind, {}).get("count", 0))
 
@@ -1651,16 +1656,21 @@ def count(
         at: Timestamp | None = None,
         branch: str | None = None,
         timeout: int | None = None,
+        **kwargs: Any,
     ) -> int:
         """Return the number of nodes of a given kind."""
+        filters = kwargs
         schema = self.schema.get(kind=kind, branch=branch)
 
         branch = branch or self.default_branch
         if at:
             at = Timestamp(at)
 
         response = self.execute_graphql(
-            query=Query(query={schema.kind: {"count": None}}).render(), branch_name=branch, at=at, timeout=timeout
+            query=Query(query={schema.kind: {"count": None, "@filters": filters}}).render(),
+            branch_name=branch,
+            at=at,
+            timeout=timeout,
         )
         return int(response.get(schema.kind, {}).get("count", 0))
 

tests/integration/test_infrahub_client.py

@@ -142,6 +142,14 @@ async def test_create_branch_async(self, client: InfrahubClient, base_dataset):
         task_id = await client.branch.create(branch_name="new-branch-2", wait_until_completion=False)
         assert isinstance(task_id, str)
 
+    async def test_count(self, client: InfrahubClient, base_dataset):
+        count = await client.count(kind=TESTING_PERSON)
+        assert count == 3
+
+    async def test_count_with_filter(self, client: InfrahubClient, base_dataset):
+        count = await client.count(kind=TESTING_PERSON, name__values=["Liam Walker", "Ethan Carter"])
+        assert count == 2
+
     # async def test_get_generic_filter_source(self, client: InfrahubClient, base_dataset):
     #     admin = await client.get(kind="CoreAccount", name__value="admin")
 

tests/unit/sdk/test_client.py

@@ -83,6 +83,16 @@ async def test_method_count(clients, mock_query_repository_count, client_type):
     assert count == 5
 
 
+@pytest.mark.parametrize("client_type", client_types)
+async def test_method_count_with_filter(clients, mock_query_repository_count, client_type):  # pylint: disable=unused-argument
+    if client_type == "standard":
+        count = await clients.standard.count(kind="CoreRepository", name__value="test")
+    else:
+        count = clients.sync.count(kind="CoreRepository", name__value="test")
+
+    assert count == 5
+
+
 @pytest.mark.parametrize("client_type", client_types)
 async def test_method_get_version(clients, mock_query_infrahub_version, client_type):  # pylint: disable=unused-argument
     if client_type == "standard":