From 1fc6b743babb5747e3edadcf92d3c917673d94a8 Mon Sep 17 00:00:00 2001
From: sadilchamishka <sandilchamishka@gmail.com>
Date: Wed, 24 Apr 2024 20:39:40 +0530
Subject: [PATCH] Fix test failure

---
 .../FederatedIdpInitLogoutProcessor.java          | 15 ++++++++++-----
 .../FederatedIdpInitLogoutProcessorTest.java      |  3 ++-
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/components/org.wso2.carbon.identity.application.authenticator.oidc/src/main/java/org/wso2/carbon/identity/application/authenticator/oidc/logout/idpinit/processor/FederatedIdpInitLogoutProcessor.java b/components/org.wso2.carbon.identity.application.authenticator.oidc/src/main/java/org/wso2/carbon/identity/application/authenticator/oidc/logout/idpinit/processor/FederatedIdpInitLogoutProcessor.java
index 78e9839c..8e24542b 100644
--- a/components/org.wso2.carbon.identity.application.authenticator.oidc/src/main/java/org/wso2/carbon/identity/application/authenticator/oidc/logout/idpinit/processor/FederatedIdpInitLogoutProcessor.java
+++ b/components/org.wso2.carbon.identity.application.authenticator.oidc/src/main/java/org/wso2/carbon/identity/application/authenticator/oidc/logout/idpinit/processor/FederatedIdpInitLogoutProcessor.java
@@ -189,11 +189,7 @@ protected LogoutResponse.LogoutResponseBuilder logoutUsingSid(String sid)
             ServerSessionManagementService serverSessionManagementService =
                     OpenIDConnectAuthenticatorDataHolder.getInstance().getServerSessionManagementService();
             serverSessionManagementService.removeSession(sessionId);
-            try {
-                UserSessionStore.getInstance().removeFederatedAuthSessionInfo(sessionId);
-            } catch (UserSessionException e) {
-                throw new LogoutServerException("Exception occurred while removing federated IDP session mapping.");
-            }
+            removeFederatedIDPSessionMapping(sessionId);
             if (log.isDebugEnabled()) {
                 log.debug("Session terminated for session Id: " + sessionId);
             }
@@ -721,4 +717,13 @@ private IdentityProvider getResidentIDPForIssuer(String tenantDomain, String jwt
         }
         return jwtIssuer.equals(issuer) ? residentIdentityProvider : null;
     }
+
+    private void removeFederatedIDPSessionMapping(String sessionID) throws LogoutServerException {
+
+        try {
+            UserSessionStore.getInstance().removeFederatedAuthSessionInfo(sessionID);
+        } catch (UserSessionException e) {
+            throw new LogoutServerException("Exception occurred while removing federated IDP session mapping.");
+        }
+    }
 }
diff --git a/components/org.wso2.carbon.identity.application.authenticator.oidc/src/test/java/org/wso2/carbon/identity/application/authenticator/oidc/logout/idpinit/processor/FederatedIdpInitLogoutProcessorTest.java b/components/org.wso2.carbon.identity.application.authenticator.oidc/src/test/java/org/wso2/carbon/identity/application/authenticator/oidc/logout/idpinit/processor/FederatedIdpInitLogoutProcessorTest.java
index 9237f859..b2fcc5a4 100644
--- a/components/org.wso2.carbon.identity.application.authenticator.oidc/src/test/java/org/wso2/carbon/identity/application/authenticator/oidc/logout/idpinit/processor/FederatedIdpInitLogoutProcessorTest.java
+++ b/components/org.wso2.carbon.identity.application.authenticator.oidc/src/test/java/org/wso2/carbon/identity/application/authenticator/oidc/logout/idpinit/processor/FederatedIdpInitLogoutProcessorTest.java
@@ -72,6 +72,7 @@
 import javax.sql.DataSource;
 import javax.xml.stream.XMLInputFactory;
 
+import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.mock;
 import static org.powermock.api.mockito.PowerMockito.mockStatic;
@@ -485,7 +486,7 @@ public void testOidcFederatedLogout() throws Exception {
         DataSource dataSource = mock(DataSource.class);
         mockStatic(IdentityDatabaseUtil.class);
         when(IdentityDatabaseUtil.getDataSource()).thenReturn(dataSource);
-        when(IdentityDatabaseUtil.getDBConnection(false)).thenReturn(getConnection(DB_NAME));
+        when(IdentityDatabaseUtil.getDBConnection(anyBoolean())).thenReturn(getConnection(DB_NAME));
         when(dataSource.getConnection()).thenReturn(getConnection(DB_NAME));
 
         // Mock the server session management service.