From 8f5eb06957cacef005f3c146be3145e59763d9a6 Mon Sep 17 00:00:00 2001 From: Quan Pham Date: Mon, 4 May 2026 14:56:55 -0400 Subject: [PATCH] Allow validation of Keycloak membership `validate_allocations` command will now add users to Keycloak groups if they are part of an allocation, and remove them if they are not. Updated keyclaok functional tests --- src/coldfront_plugin_cloud/kc_client.py | 19 +++++ .../commands/validate_allocations.py | 67 +++++++++++++++ .../functional/keycloak/test_keycloak.py | 81 +++++++++++++++---- 3 files changed, 150 insertions(+), 17 deletions(-) diff --git a/src/coldfront_plugin_cloud/kc_client.py b/src/coldfront_plugin_cloud/kc_client.py index d4ed83de..96a8b06f 100644 --- a/src/coldfront_plugin_cloud/kc_client.py +++ b/src/coldfront_plugin_cloud/kc_client.py @@ -109,3 +109,22 @@ def get_user_groups(self, user_id) -> list[str]: r.raise_for_status() groups = GroupResponse.model_validate(r.json()) return [group.name for group in groups.root] + + def get_group_members(self, group_id) -> list[str]: + url = f"{self.base_url}/admin/realms/{self.realm}/groups/{group_id}/members" + page_size = 100 # Default KeyCloak page size https://www.keycloak.org/docs-api/latest/rest-api/index.html#_query_parameters_32 + page_offset = 0 + users = [] + + while True: + r = self.api_client.get( + url, params={"first": page_offset, "max": page_size} + ) + r.raise_for_status() + batch = UserResponse.model_validate(r.json()) + users.extend(user.username for user in batch.root) + if len(batch.root) < page_size: + break + page_offset += page_size + + return users diff --git a/src/coldfront_plugin_cloud/management/commands/validate_allocations.py b/src/coldfront_plugin_cloud/management/commands/validate_allocations.py index 1edf04d9..96097632 100644 --- a/src/coldfront_plugin_cloud/management/commands/validate_allocations.py +++ b/src/coldfront_plugin_cloud/management/commands/validate_allocations.py @@ -3,11 +3,13 @@ from coldfront_plugin_cloud import attributes from coldfront_plugin_cloud import utils from coldfront_plugin_cloud import tasks +from coldfront_plugin_cloud import signals from django.core.management.base import BaseCommand from coldfront.core.resource.models import Resource from coldfront.core.allocation.models import ( Allocation, + AllocationUser, ) from keystoneauth1.exceptions import http from kubernetes.dynamic import exceptions as k8s_exceptions @@ -46,6 +48,66 @@ def check_institution_specific_code(self, allocation, apply): utils.set_attribute_on_allocation(allocation, attr, "N/A") logger.warning(f'Attribute "{attr}" added to allocation {alloc_str}') + def validate_keycloak_group_memberships(self, allocation: Allocation, apply: bool): + # Fetch or cache Keycloak client + kc_client = getattr(self, "_kc_client", None) + if kc_client is None: + kc_client = tasks.get_kc_client() + self._kc_client = kc_client + + resource = allocation.resources.first() + group_template = resource.get_attribute( + attributes.RESOURCE_KEYCLOAK_GROUP_TEMPLATE + ) + if not group_template: + logger.info( + f"Keycloak enabled but no group name template specified for resource {resource.name}. Skipping validation" + ) + return + + allocation_users: list[AllocationUser] = ( + allocation.allocationuser_set.select_related("user").all() + ) + allocation_usernames: set[str] = set( + au.user.username for au in allocation_users + ) + + group_name = tasks._get_keycloak_group_name(allocation, group_template) + + try: + group_name = tasks._get_keycloak_group_name(allocation, group_template) + except KeyError as e: + logger.error( + f"Invalid Keycloak group template for allocation {allocation.pk}: missing template variable {e}. Skipping validation" + ) + return + + group_id = kc_client.get_group_id(group_name) + if group_id: + group_usernames = set(kc_client.get_group_members(group_id)) + else: + group_usernames = set() + + to_add = [ + au for au in allocation_users if au.user.username not in group_usernames + ] + to_remove = group_usernames - allocation_usernames + + for au in to_add: + username = au.user.username + logger.info(f"Adding user {username} to Keycloak group {group_name}") + if apply: + tasks.add_user_to_keycloak(au.pk) + for username in to_remove: + logger.info(f"Removing user {username} from Keycloak group {group_name}") + if apply: + if user_id := kc_client.get_user_id(username): + kc_client.remove_user_from_group(user_id, group_id) + else: + logger.warning( + f"User {username} not found in Keycloak, cannot remove from group {group_name}." + ) + def handle(self, *args, **options): for resource_name in self.PLUGIN_RESOURCE_NAMES: resource = Resource.objects.filter(resource_type__name=resource_name) @@ -70,6 +132,11 @@ def handle(self, *args, **options): ) continue + if signals.is_keycloak_enabled(): + self.validate_keycloak_group_memberships( + allocation, options["apply"] + ) + # Check project exists in remote cluster try: allocator.get_project(project_id) diff --git a/src/coldfront_plugin_cloud/tests/functional/keycloak/test_keycloak.py b/src/coldfront_plugin_cloud/tests/functional/keycloak/test_keycloak.py index e515d82a..06ba22d0 100644 --- a/src/coldfront_plugin_cloud/tests/functional/keycloak/test_keycloak.py +++ b/src/coldfront_plugin_cloud/tests/functional/keycloak/test_keycloak.py @@ -1,4 +1,7 @@ +from unittest import mock + from django.contrib.auth.models import User +from django.core.management import call_command from coldfront.core.resource.models import ResourceAttribute, ResourceAttributeType from coldfront_plugin_cloud import tasks, kc_client, attributes, utils @@ -21,6 +24,17 @@ def setUpTestData(cls) -> None: value="$resource_name/$allocated_project_id", ) + def setUp(self) -> None: + mock_allocator = mock.MagicMock() + mock_allocator.allocation_str = "Test Allocation of project Test Project" + self.patcher = mock.patch( + "coldfront_plugin_cloud.tasks.find_allocator", return_value=mock_allocator + ) + self.mock_find_allocator = self.patcher.start() + + def tearDown(self) -> None: + self.patcher.stop() + def new_keycloak_user(self, cf_username): url = f"{self.kc_admin_client.base_url}/admin/realms/{self.kc_admin_client.realm}/users" payload = { @@ -51,10 +65,10 @@ def test_user_added_to_allocation(self): user = self.new_user() project = self.new_project(pi=user) allocation = self.new_allocation(project, self.resource, 1) - allocation_user = self.new_allocation_user(allocation, user) + self.new_allocation_user(allocation, user) - # Simulate triggering the allocation activate signal - tasks.add_user_to_keycloak(allocation_user.pk) + # Validation should add user to Keycloak group + call_command("validate_allocations", apply=True) # Check that the user exists in Keycloak user_id = self.kc_admin_client.get_user_id(user.username) @@ -72,13 +86,14 @@ def test_user_removed_from_allocation(self): allocation = self.new_allocation(project, self.resource, 1) allocation_user = self.new_allocation_user(allocation, user) - tasks.add_user_to_keycloak(allocation_user.pk) + call_command("validate_allocations", apply=True) user_id = self.kc_admin_client.get_user_id(user.username) user_groups = self.kc_admin_client.get_user_groups(user_id) self.assertIn(f"{self.resource.name}/Test Value", user_groups) - tasks.remove_user_from_keycloak(allocation_user.pk) + allocation_user.delete() + call_command("validate_allocations", apply=True) # Check that the user is no longer in the group user_groups = self.kc_admin_client.get_user_groups(user_id) @@ -91,10 +106,10 @@ def test_user_not_in_keycloak_added_to_allocation(self): allocation = self.new_allocation( project, self.resource, 1, attr_value="Test Not Created" ) - allocation_user = self.new_allocation_user(allocation, user) + self.new_allocation_user(allocation, user) # Should not raise error - tasks.add_user_to_keycloak(allocation_user.pk) + call_command("validate_allocations", apply=True) user_id = self.kc_admin_client.get_user_id(user.username) self.assertIsNone(user_id) @@ -127,12 +142,10 @@ def test_multiple_users_in_same_allocation(self): # Add multiple users to the allocation users = [self.new_user() for _ in range(3)] - allocation_users = [ - self.new_allocation_user(allocation, user) for user in users - ] + for user in users: + self.new_allocation_user(allocation, user) - for allocation_user in allocation_users: - tasks.add_user_to_keycloak(allocation_user.pk) + call_command("validate_allocations", apply=True) # Verify all users are in the group for user in users: @@ -151,8 +164,7 @@ def test_remove_one_user_keeps_others_in_group(self): self.new_allocation_user(allocation, user) for user in users ] - for allocation_user in allocation_users: - tasks.add_user_to_keycloak(allocation_user.pk) + call_command("validate_allocations", apply=True) tasks.remove_user_from_keycloak(allocation_users[0].pk) @@ -181,10 +193,9 @@ def test_user_in_multiple_allocations_groups(self): # Add user to both allocations allocation_user1 = self.new_allocation_user(allocation1, user) - allocation_user2 = self.new_allocation_user(allocation2, user) + self.new_allocation_user(allocation2, user) - tasks.add_user_to_keycloak(allocation_user1.pk) - tasks.add_user_to_keycloak(allocation_user2.pk) + call_command("validate_allocations", apply=True) # Verify user is in both groups user_id = self.kc_admin_client.get_user_id(user.username) @@ -229,3 +240,39 @@ def test_user_added_without_keycloak_group_template(self): self.assertIsNotNone(user_id) user_groups = self.kc_admin_client.get_user_groups(user_id) self.assertEqual(user_groups, []) + + +class TestKeyCloakGetGroupMembersPagination(base.TestBase): + @classmethod + def setUpTestData(cls) -> None: + super().setUpTestData() + cls.kc_admin_client = kc_client.KeyCloakAPIClient() + + def new_keycloak_user(self, cf_username): + url = f"{self.kc_admin_client.base_url}/admin/realms/{self.kc_admin_client.realm}/users" + payload = { + "username": cf_username, + "enabled": True, + "email": cf_username, + } + r = self.kc_admin_client.api_client.post(url, json=payload) + r.raise_for_status() + + def test_get_group_members_pagination(self): + group_name = "Test Pagination Group" + self.kc_admin_client.create_group(group_name) + group_id = self.kc_admin_client.get_group_id(group_name) + + # Create 250 users and add them to the group (to ensure pagination is needed, as page size is 100) + for i in range(250): + username = f"pagination_user_{i}@example.com" + self.new_user(username=username) + self.new_keycloak_user(username) + self.kc_admin_client.add_user_to_group( + self.kc_admin_client.get_user_id(username), group_id + ) + + members = self.kc_admin_client.get_group_members(group_id) + self.assertEqual(len(members), 250) + expected_usernames = {f"pagination_user_{i}@example.com" for i in range(250)} + self.assertEqual(set(members), expected_usernames)