From 329e902b46034f3561449cc0e49be5caf4969628 Mon Sep 17 00:00:00 2001
From: viniciusdc <vinivdc2009@hotmail.com>
Date: Tue, 24 Sep 2024 19:01:02 -0300
Subject: [PATCH] fix failing ci test

---
 src/_nebari/keycloak.py          | 41 +++++++++++++++++++-------------
 src/_nebari/upgrade.py           | 28 +++++++++++++++-------
 tests/tests_unit/test_upgrade.py |  5 ++++
 3 files changed, 50 insertions(+), 24 deletions(-)

diff --git a/src/_nebari/keycloak.py b/src/_nebari/keycloak.py
index ea8815940d..6bfea9b8b3 100644
--- a/src/_nebari/keycloak.py
+++ b/src/_nebari/keycloak.py
@@ -81,27 +81,16 @@ def list_users(keycloak_admin: keycloak.KeycloakAdmin):
         )
 
 
-def get_keycloak_admin_from_config(config: schema.Main):
-    keycloak_server_url = os.environ.get(
-        "KEYCLOAK_SERVER_URL", f"https://{config.domain}/auth/"
-    )
-
-    keycloak_username = os.environ.get("KEYCLOAK_ADMIN_USERNAME", "root")
-    keycloak_password = os.environ.get(
-        "KEYCLOAK_ADMIN_PASSWORD", config.security.keycloak.initial_root_password
-    )
-
-    should_verify_tls = config.certificate.type != CertificateEnum.selfsigned
-
+def get_keycloak_admin(server_url, username, password, verify=False):
     try:
         keycloak_admin = keycloak.KeycloakAdmin(
-            server_url=keycloak_server_url,
-            username=keycloak_username,
-            password=keycloak_password,
+            server_url=server_url,
+            username=username,
+            password=password,
             realm_name=os.environ.get("KEYCLOAK_REALM", "nebari"),
             user_realm_name="master",
             auto_refresh_token=("get", "put", "post", "delete"),
-            verify=should_verify_tls,
+            verify=verify,
         )
     except (
         keycloak.exceptions.KeycloakConnectionError,
@@ -112,6 +101,26 @@ def get_keycloak_admin_from_config(config: schema.Main):
     return keycloak_admin
 
 
+def get_keycloak_admin_from_config(config: schema.Main):
+    keycloak_server_url = os.environ.get(
+        "KEYCLOAK_SERVER_URL", f"https://{config.domain}/auth/"
+    )
+
+    keycloak_username = os.environ.get("KEYCLOAK_ADMIN_USERNAME", "root")
+    keycloak_password = os.environ.get(
+        "KEYCLOAK_ADMIN_PASSWORD", config.security.keycloak.initial_root_password
+    )
+
+    should_verify_tls = config.certificate.type != CertificateEnum.selfsigned
+
+    return get_keycloak_admin(
+        server_url=keycloak_server_url,
+        username=keycloak_username,
+        password=keycloak_password,
+        verify=should_verify_tls,
+    )
+
+
 def keycloak_rest_api_call(config: schema.Main = None, request: str = None):
     """Communicate directly with the Keycloak REST API by passing it a request"""
     keycloak_server_url = os.environ.get(
diff --git a/src/_nebari/upgrade.py b/src/_nebari/upgrade.py
index 56cb565f0d..bcc08be0f5 100644
--- a/src/_nebari/upgrade.py
+++ b/src/_nebari/upgrade.py
@@ -24,7 +24,7 @@
 from typing_extensions import override
 
 from _nebari.config import backup_configuration
-from _nebari.keycloak import get_keycloak_admin_from_config
+from _nebari.keycloak import get_keycloak_admin
 from _nebari.stages.infrastructure import (
     provider_enum_default_node_groups_map,
     provider_enum_name_map,
@@ -1256,21 +1256,33 @@ def _version_specific_upgrade(
         rich.print(text)
 
         confirm = Prompt.ask(
-            "[bold]Would you like Nebari to update your group permissions now?[/bold] (y/n)",
+            "[bold]Would you like Nebari to update your group permissions now?[/bold]",
             choices=["y", "N"],
             default="N",
         )
-
         if confirm.lower() == "y":
             # Proceed with updating group permissions
-            keycloak_admin = get_keycloak_admin_from_config(config)
+            keycloak_admin = get_keycloak_admin(
+                server_url=f"https://{config['domain']}/auth/",
+                username="root",
+                password=config["security"]["keycloak"]["initial_root_password"],
+            )
+            client_id = keycloak_admin.get_client_id("jupyterhub")
+            _role_representation = keycloak_admin.get_role_by_id(
+                role_id=keycloak_admin.get_client_role_id(
+                    client_id=client_id, role_name="allow-group-directory-creation-role"
+                )
+            )
             groups = keycloak_admin.get_groups()
+            groups_with_roles = keycloak_admin.get_client_role_groups(
+                client_id=client_id, role_name="allow-group-directory-creation-role"
+            )
             groups_without_role = [
                 group
                 for group in groups
-                if "allow-group-directory-creation-role"
-                not in group.get("attributes", {})
+                if group["id"] not in [group["id"] for group in groups_with_roles]
             ]
+
             if groups_without_role:
                 group_names = ", ".join(
                     [group["name"] for group in groups_without_role]
@@ -1282,8 +1294,8 @@ def _version_specific_upgrade(
                     _group_id = group["id"]
                     keycloak_admin.assign_group_client_roles(
                         group_id=_group_id,
-                        client_id="jupyterhub",
-                        roles=["allow-group-directory-creation-role"],
+                        client_id=client_id,
+                        roles=[_role_representation],
                     )
                 rich.print(
                     "[green]Group permissions have been updated successfully.[/green]"
diff --git a/tests/tests_unit/test_upgrade.py b/tests/tests_unit/test_upgrade.py
index a19095726b..d328c74e9f 100644
--- a/tests/tests_unit/test_upgrade.py
+++ b/tests/tests_unit/test_upgrade.py
@@ -67,6 +67,11 @@ def mock_input(prompt, **kwargs):
             == "Have you backed up your custom dashboards (if necessary), deleted the prometheus-node-exporter daemonset and updated the kube-prometheus-stack CRDs?"
         ):
             return "y"
+        elif (
+            prompt
+            == "[bold]Would you like Nebari to update your group permissions now?[/bold]"
+        ):
+            return "N"
         # All other prompts will be answered with "y"
         else:
             return "y"