Fix NPE in Organization Audience Shared Role Deletion

components/org.wso2.carbon.identity.organization.management.handler/src/main/java/org/wso2/carbon/identity/organization/management/handler/listener/SharedRoleMgtListener.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023-2025, WSO2 LLC. (http://www.wso2.com).
+ * Copyright (c) (2023-2025), WSO2 LLC. (https://www.wso2.com).
  *
  * WSO2 LLC. licenses this file to you under the Apache License,
  * Version 2.0 (the "License"); you may not use this file except
@@ -326,9 +326,8 @@ private void handleRemovedOrganizationAudienceRolesOnAppUpdate(List<RoleV2> remo
             CompletableFuture.runAsync(() -> {
                 String sharedAppOrgId = sharedApplication.getOrganizationId();
                 try {
-                    handleOrganizationAudiencedSharedRoleDeletion(removedOrgRolesList,
-                            serviceProvider.getApplicationResourceId(),
-                            tenantDomain, sharedAppOrgId);
+                    handleOrganizationAudiencedSharedRoleDeletion(
+                            removedOrgRolesList, mainAppId, tenantDomain, sharedAppOrgId);
                 } catch (IdentityRoleManagementException | OrganizationManagementException e) {
                     LOG.error(String.format("Exception occurred during deleting roles from organization %s",
                             sharedApplication.getOrganizationId()), e);
@@ -536,9 +535,16 @@ private void handleOrganizationAudiencedSharedRoleDeletion(List<RoleV2> rolesLis
 
         // Get each role associated applications.
         for (String mainAppRoleId : mainAppRoleIds) {
-            List<String> associatedApplicationsIds =
-                    roleManagementService.getAssociatedApplicationByRoleId(mainAppRoleId,
-                            mainApplicationTenantDomain);
+            List<String> associatedApplicationsIds;
+            try {
+                PrivilegedCarbonContext.startTenantFlow();
+                PrivilegedCarbonContext.getThreadLocalCarbonContext()
+                        .setTenantDomain(mainApplicationTenantDomain, true);
+                associatedApplicationsIds = roleManagementService.getAssociatedApplicationByRoleId(mainAppRoleId,
+                        mainApplicationTenantDomain);
+            } finally {
+                PrivilegedCarbonContext.endTenantFlow();
+            }
             String sharedRoleId = mainRoleToSharedRoleMappingsInSubOrg.get(mainAppRoleId);
             if (StringUtils.isBlank(sharedRoleId)) {
                 // There is no role available in the shared org. May be due to role creation issue.

components/org.wso2.carbon.identity.organization.management.handler/src/test/java/org/wso2/carbon/identity/organization/management/handler/SharedRoleMgtHandlerTest.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024-2025, WSO2 LLC. (http://www.wso2.com).
+ * Copyright (c) (2024-2025), WSO2 LLC. (https://www.wso2.com).
  *
  * WSO2 LLC. licenses this file to you under the Apache License,
  * Version 2.0 (the "License"); you may not use this file except
@@ -22,7 +22,6 @@
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
-import org.wso2.carbon.base.CarbonBaseConstants;
 import org.wso2.carbon.context.PrivilegedCarbonContext;
 import org.wso2.carbon.identity.application.common.model.RoleV2;
 import org.wso2.carbon.identity.application.mgt.ApplicationManagementService;
@@ -32,10 +31,10 @@
 import org.wso2.carbon.identity.event.event.Event;
 import org.wso2.carbon.identity.organization.management.application.constant.OrgApplicationMgtConstants;
 import org.wso2.carbon.identity.organization.management.handler.internal.OrganizationManagementHandlerDataHolder;
+import org.wso2.carbon.identity.organization.management.handler.util.TestUtils;
 import org.wso2.carbon.identity.organization.management.service.OrganizationManager;
 import org.wso2.carbon.identity.role.v2.mgt.core.RoleManagementService;
 
-import java.nio.file.Paths;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
@@ -71,7 +70,7 @@ public class SharedRoleMgtHandlerTest {
     @BeforeClass
     public void setUp() {
 
-        initPrivilegedCarbonContext();
+        TestUtils.initPrivilegedCarbonContext();
         PrivilegedCarbonContext.getThreadLocalCarbonContext().setTenantDomain(PARENT_ORG_TENANT_DOMAIN);
         PrivilegedCarbonContext.getThreadLocalCarbonContext().setUsername(PARENT_ORG_USER_NAME);
 
@@ -193,15 +192,6 @@ public void testHandleEventForPreShareApplicationEventWithApplicationAud() throw
         sharedRoleMgtHandler.handleEvent(event);
     }
 
-    private void initPrivilegedCarbonContext() {
-
-        System.setProperty(
-                CarbonBaseConstants.CARBON_HOME,
-                Paths.get(System.getProperty("user.dir"), "src", "test", "resources").toString()
-        );
-        PrivilegedCarbonContext.startTenantFlow();
-    }
-
     private static Event createPreShareApplicationEvent() {
 
         Event event = new Event(OrgApplicationMgtConstants.EVENT_PRE_SHARE_APPLICATION);

components/org.wso2.carbon.identity.organization.management.handler/src/test/java/org/wso2/carbon/identity/organization/management/handler/listener/SharedRoleMgtListenerTest.java

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2025, WSO2 LLC. (http://www.wso2.com).
+ * Copyright (c) 2025, WSO2 LLC. (https://www.wso2.com).
  *
  * WSO2 LLC. licenses this file to you under the Apache License,
  * Version 2.0 (the "License"); you may not use this file except
@@ -22,6 +22,7 @@
 import org.mockito.Mock;
 import org.mockito.MockedStatic;
 import org.mockito.MockitoAnnotations;
+import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.DataProvider;
 import org.testng.annotations.Test;
@@ -31,10 +32,19 @@
 import org.wso2.carbon.identity.application.common.model.ServiceProviderProperty;
 import org.wso2.carbon.identity.application.mgt.ApplicationManagementService;
 import org.wso2.carbon.identity.core.util.IdentityTenantUtil;
+import org.wso2.carbon.identity.core.util.IdentityUtil;
 import org.wso2.carbon.identity.organization.management.application.OrgApplicationManager;
+import org.wso2.carbon.identity.organization.management.application.model.SharedApplication;
 import org.wso2.carbon.identity.organization.management.handler.internal.OrganizationManagementHandlerDataHolder;
+import org.wso2.carbon.identity.organization.management.handler.util.TestUtils;
 import org.wso2.carbon.identity.organization.management.service.OrganizationManager;
 import org.wso2.carbon.identity.organization.management.service.util.OrganizationManagementUtil;
+import org.wso2.carbon.identity.role.v2.mgt.core.RoleManagementService;
+
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
 
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.mockStatic;
@@ -55,9 +65,13 @@ public class SharedRoleMgtListenerTest {
     private static final String SAMPLE_MAIN_APP_ID = "main-app-id";
     private static final String SAMPLE_SHARED_APP_ID = "shared-app-id";
     private static final String SAMPLE_SHARED_APP_ORG_ID = "shared-app-org-id";
+    private static final String SAMPLE_ORG_ID = "org-id";
+    private static final String SAMPLE_ROLE_ID = "role-id";
+    private static final String SAMPLE_SHARED_ROLE_ID = "shared-role-id";
     private static final String IS_FRAGMENT_APP = "isFragmentApp";
     private static final String ORGANIZATION_AUD = "organization";
     private static final String APPLICATION_AUD = "application";
+    private static final String REMOVED_ORGANIZATION_AUDIENCE_ROLES = "removedOrganizationAudienceRoles";
 
     @Mock
     private ApplicationManagementService mockedApplicationManagementService;
@@ -68,21 +82,33 @@ public class SharedRoleMgtListenerTest {
     @Mock
     private OrgApplicationManager mockedOrgApplicationManager;
 
+    @Mock
+    private RoleManagementService mockedRoleManagementService;
+
     private MockedStatic<OrganizationManagementUtil> organizationManagementUtilMockedStatic;
     private MockedStatic<IdentityTenantUtil> identityTenantUtilMockedStatic;
 
     @BeforeClass
     public void setUpClass() {
 
+        TestUtils.initPrivilegedCarbonContext();
         MockitoAnnotations.openMocks(this);
         OrganizationManagementHandlerDataHolder.getInstance().
                 setApplicationManagementService(mockedApplicationManagementService);
         OrganizationManagementHandlerDataHolder.getInstance().setOrganizationManager(mockedOrganizationManager);
         OrganizationManagementHandlerDataHolder.getInstance().setOrgApplicationManager(mockedOrgApplicationManager);
+        OrganizationManagementHandlerDataHolder.getInstance().setRoleManagementServiceV2(mockedRoleManagementService);
         organizationManagementUtilMockedStatic = mockStatic(OrganizationManagementUtil.class);
         identityTenantUtilMockedStatic = mockStatic(IdentityTenantUtil.class);
     }
 
+    @AfterClass
+    public void tearDown() {
+
+        organizationManagementUtilMockedStatic.close();
+        identityTenantUtilMockedStatic.close();
+    }
+
     @DataProvider(name = "organizationTypeDataProvider")
     public Object[][] organizationTypeDataProvider() {
 
@@ -183,4 +209,35 @@ public void testDoPostGetAllowedAudienceForRoleAssociation(String mainAppId) thr
             assertEquals(associatedRolesConfig.getAllowedAudience(), ORGANIZATION_AUD);
         }
     }
+
+    @Test
+    public void testHandleRemovedOrganizationAudienceRolesOnAppUpdate() throws Exception {
+
+        RoleV2 roleV2 = new RoleV2(SAMPLE_ROLE_ID, SAMPLE_ROLE_NAME);
+        List<RoleV2> removedOrgRolesList = Collections.singletonList(roleV2);
+
+        Map<String, Object> threadLocalProperties = new HashMap<>();
+        threadLocalProperties.put(REMOVED_ORGANIZATION_AUDIENCE_ROLES, removedOrgRolesList);
+        IdentityUtil.threadLocalProperties.set(threadLocalProperties);
+
+        SharedApplication sharedApplication = new SharedApplication(SAMPLE_SHARED_APP_ID, SAMPLE_ORG_ID);
+        List<SharedApplication> sharedApplications = Collections.singletonList(sharedApplication);
+
+        Map<String, String> mainRoleToSharedRoleMappingsInSubOrg = new HashMap<>();
+        mainRoleToSharedRoleMappingsInSubOrg.put(SAMPLE_ROLE_ID, SAMPLE_SHARED_ROLE_ID);
+
+        when(mockedOrgApplicationManager.getSharedApplications(null, SAMPLE_SHARED_APP_ID))
+                .thenReturn(sharedApplications);
+        when(mockedRoleManagementService.getMainRoleToSharedRoleMappingsBySubOrg(
+                Collections.singletonList(SAMPLE_ROLE_ID), null))
+                .thenReturn(mainRoleToSharedRoleMappingsInSubOrg);
+        when(mockedRoleManagementService.getAssociatedApplicationByRoleId(SAMPLE_ROLE_ID, SAMPLE_TENANT_DOMAIN))
+                .thenReturn(Collections.singletonList(SAMPLE_SHARED_APP_ID));
+
+        SharedRoleMgtListener sharedRoleMgtListener = new SharedRoleMgtListener();
+        ServiceProvider serviceProvider = new ServiceProvider();
+        serviceProvider.setApplicationResourceId(SAMPLE_SHARED_APP_ID);
+        assertEquals(sharedRoleMgtListener.doPostUpdateApplication(serviceProvider, SAMPLE_TENANT_DOMAIN,
+                SAMPLE_USERNAME), true);
+    }
 }

components/org.wso2.carbon.identity.organization.management.handler/src/test/java/org/wso2/carbon/identity/organization/management/handler/util/TestUtils.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright (c) 2025, WSO2 LLC. (https://www.wso2.com).
+ *
+ * WSO2 LLC. licenses this file to you under the Apache License,
+ * Version 2.0 (the "License"); you may not use this file except
+ * in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.wso2.carbon.identity.organization.management.handler.util;
+
+import org.wso2.carbon.base.CarbonBaseConstants;
+import org.wso2.carbon.context.PrivilegedCarbonContext;
+
+import java.nio.file.Paths;
+
+/**
+ * Test utility class for the organization management handler.
+ */
+public class TestUtils {
+
+    public static void initPrivilegedCarbonContext() {
+
+        System.setProperty(
+                CarbonBaseConstants.CARBON_HOME,
+                Paths.get(System.getProperty("user.dir"), "src", "test", "resources").toString()
+                          );
+        PrivilegedCarbonContext.startTenantFlow();
+    }
+}