From 00521ea42c0a531c3ebdecfb08467058a48c070f Mon Sep 17 00:00:00 2001
From: Ana Paula Mendes <anapaula.ds.mendes@gmail.com>
Date: Fri, 23 Jun 2023 02:17:05 -0300
Subject: [PATCH] feat: Add query to get users group method and permit
 pagination (#444)

* fix: Add query to get users group method and permit pagination

* test: Add test to user groups pagination

* fix: Add blank line
---
 src/keycloak/keycloak_admin.py | 20 +++++++++++++++-----
 tests/test_keycloak_admin.py   | 29 +++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+), 5 deletions(-)

diff --git a/src/keycloak/keycloak_admin.py b/src/keycloak/keycloak_admin.py
index 9c80057..20e5883 100644
--- a/src/keycloak/keycloak_admin.py
+++ b/src/keycloak/keycloak_admin.py
@@ -867,24 +867,34 @@ class KeycloakAdmin:
         data_raw = self.connection.raw_get(urls_patterns.URL_ADMIN_USER.format(**params_path))
         return raise_error_from_response(data_raw, KeycloakGetError)
 
-    def get_user_groups(self, user_id, brief_representation=True):
+    def get_user_groups(self, user_id, query=None, brief_representation=True):
         """Get user groups.
 
         Returns a list of groups of which the user is a member
 
         :param user_id: User id
         :type user_id: str
+        :param query: Additional query options
+        :type query: dict
         :param brief_representation: whether to omit attributes in the response
         :type brief_representation: bool
         :return: user groups list
         :rtype: list
         """
+        query = query or {}
+
         params = {"briefRepresentation": brief_representation}
+
+        query.update(params)
+
         params_path = {"realm-name": self.connection.realm_name, "id": user_id}
-        data_raw = self.connection.raw_get(
-            urls_patterns.URL_ADMIN_USER_GROUPS.format(**params_path), **params
-        )
-        return raise_error_from_response(data_raw, KeycloakGetError)
+
+        url = urls_patterns.URL_ADMIN_USER_GROUPS.format(**params_path)
+
+        if "first" in query or "max" in query:
+            return self.__fetch_paginated(url, query)
+
+        return self.__fetch_all(url, query)
 
     def update_user(self, user_id, payload):
         """Update the user.
diff --git a/tests/test_keycloak_admin.py b/tests/test_keycloak_admin.py
index 454acfc..8be5229 100644
--- a/tests/test_keycloak_admin.py
+++ b/tests/test_keycloak_admin.py
@@ -338,6 +338,35 @@ def test_users_pagination(admin: KeycloakAdmin, realm: str):
     assert len(users) == 20, len(users)
 
 
+def test_user_groups_pagination(admin: KeycloakAdmin, realm: str):
+    """Test user groups pagination.
+
+    :param admin: Keycloak Admin client
+    :type admin: KeycloakAdmin
+    :param realm: Keycloak realm
+    :type realm: str
+    """
+    admin.realm_name = realm
+
+    user_id = admin.create_user(
+        payload={"username": "username_1", "email": "username_1@test.test"}
+    )
+
+    for ind in range(admin.PAGE_SIZE + 50):
+        group_name = f"group_{ind}"
+        group_id = admin.create_group(payload={"name": group_name})
+        admin.group_user_add(user_id=user_id, group_id=group_id)
+
+    groups = admin.get_user_groups(user_id=user_id)
+    assert len(groups) == admin.PAGE_SIZE + 50, len(groups)
+
+    groups = admin.get_user_groups(user_id=user_id, query={"first": 100})
+    assert len(groups) == 50, len(groups)
+
+    groups = admin.get_user_groups(user_id=user_id, query={"max": 20})
+    assert len(groups) == 20, len(groups)
+
+
 def test_idps(admin: KeycloakAdmin, realm: str):
     """Test IDPs.