Preserve custom attributes on Discord Access Okta Group UI updates (#249)

Author: jonathanhle

api/services/okta_service.py

@@ -164,10 +164,38 @@ async def _create_group(name: str, description: str) -> Group:
     def update_group(self, groupId: str, name: str, description: str) -> Group:
         async def _update_group(groupId: str, name: str, description: str) -> Group:
             async with self._get_sessioned_okta_request_executor() as _:
+                # Fetch Existing Group Data
+                existing_group_data, _, get_error = await OktaService._retry(self.okta_client.get_group, groupId)
+                if get_error is not None:
+                    logger.error(f"Failed to fetch existing group {groupId} before update: {get_error}")
+                    raise Exception(f"Failed to fetch existing group {groupId} before update: {get_error}")
+                if existing_group_data is None:
+                    logger.error(f"Group {groupId} not found in Okta before update.")
+                    raise Exception(f"Group {groupId} not found in Okta before update.")
+
+                # Extract Existing Profile
+                existing_profile = {}
+                if existing_group_data.profile:
+                    # Using __dict__ can be fragile if Okta changes internal representation.
+                    # If specific profile attributes are known, accessing them directly might be safer.
+                    # Filter out None values if needed by Okta API or desired.
+                    existing_profile = {k: v for k, v in existing_group_data.profile.__dict__.items() if v is not None}
+
+                # Merge Updated Profile Data
+                new_profile = {**existing_profile}  # Start with a copy of the existing profile
+                new_profile["name"] = name  # Update/set the name
+                new_profile["description"] = (
+                    description if description is not None else ""
+                )  # Update/set the description (handle None)
+
+                # Construct the New Payload
+                group_payload = OktaGroupType({"profile": new_profile})
+
+                # Modify the Update Call to use the new payload
                 group, _, error = await OktaService._retry(
                     self.okta_client.update_group,
                     groupId,
-                    OktaGroupType({"profile": {"name": name, "description": description}}),
+                    group_payload,
                 )
 
             if error is not None:

tests/factories.py

@@ -29,6 +29,7 @@ class GroupProfileFactory(factory.Factory[GroupProfile]):
         {
             "name": factory.Faker("pystr"),
             "description": factory.Faker("pystr"),
+            "allow_discord_access": None,  # Default to None for the custom attribute
         }
     )
 

tests/test_okta_service.py

@@ -1,8 +1,10 @@
-from unittest.mock import MagicMock, patch
+import asyncio
+from typing import Any
+from unittest.mock import AsyncMock, MagicMock, patch
 
 from okta.models.group_rule import GroupRule as OktaGroupRuleType
 
-from api.services.okta_service import is_managed_group
+from api.services.okta_service import OktaService, is_managed_group
 
 
 def test_is_managed_group_with_allow_discord_access_false() -> None:
@@ -60,3 +62,64 @@ def test_is_managed_group_with_allow_discord_access_undefined() -> None:
         # Call the function and assert the expected result
         result = is_managed_group(group, group_ids_with_group_rules, OKTA_GROUP_PROFILE_CUSTOM_ATTR)
         assert result is True
+
+
+def test_update_group_preserves_custom_attributes() -> None:
+    """Test that update_group preserves custom attributes when updating a group."""
+    # Create a new event loop for this test
+    loop = asyncio.new_event_loop()
+    asyncio.set_event_loop(loop)
+
+    # Mock asyncio.run to use our event loop
+    with patch("asyncio.run", side_effect=lambda x: loop.run_until_complete(x)):
+        # Create OktaService instance with mock client
+        service = OktaService()
+        service.okta_client = MagicMock()
+
+        # Set up the mocks for the existing group and the update call
+        group_id = "test-group-id"
+
+        # Create a mock group with a profile that has the custom attribute
+        existing_group = MagicMock()
+        # Instead of setting __dict__ directly, configure the mock properly
+        existing_group.profile = MagicMock()
+        existing_group.profile.name = "Old Name"
+        existing_group.profile.description = "Old Description"
+        existing_group.profile.allow_discord_access = True
+
+        # Mock the get_group and update_group methods
+        service.okta_client.get_group = AsyncMock(return_value=(existing_group, None, None))
+        service.okta_client.update_group = AsyncMock(return_value=(MagicMock(), None, None))
+
+        # Create a mock for the SessionedOktaRequestExecutor context manager
+        # This is a special class that implements the async context manager protocol
+        class MockSessionedExecutor:
+            """Mock class for SessionedOktaRequestExecutor with async context manager methods"""
+
+            async def __aenter__(self) -> None:
+                """Mock for __aenter__ - called when entering an 'async with' block"""
+                return None
+
+            async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+                """Mock for __aexit__ - called when exiting an 'async with' block"""
+                return None
+
+        # Create a mock context manager instance
+        mock_executor = MockSessionedExecutor()
+
+        # Use patch to mock the _get_sessioned_okta_request_executor method
+        # This avoids directly assigning to the method, which mypy doesn't like
+        with patch.object(service, "_get_sessioned_okta_request_executor", return_value=mock_executor):
+            # Call update_group
+            service.update_group(group_id, "New Name", "New Description")
+
+            # Verify update_group was called with a payload that preserved the custom attribute
+            args, _ = service.okta_client.update_group.call_args
+            assert len(args) == 2
+            assert args[0] == group_id
+
+            # Check that the payload contains both the updated fields and the preserved custom attribute
+            updated_payload = args[1]
+            assert updated_payload.profile.name == "New Name"
+            assert updated_payload.profile.description == "New Description"
+            assert updated_payload.profile.allow_discord_access is True