Change code retrieval method

wso2 · Nov 22, 2024 · 4bc8e75 · 4bc8e75
1 parent 0f60df4
commit 4bc8e75
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 4 deletions.
diff --git a/...src/test/java/org/wso2/identity/integration/test/auth/PasswordlessSMSOTPAuthTestCase.java b/...src/test/java/org/wso2/identity/integration/test/auth/PasswordlessSMSOTPAuthTestCase.java
@@ -210,8 +210,9 @@ private void performUserLogin() throws Exception {
 
         sendLoginPostForIdentifier(client, sessionDataKey, userObject.getUserName());
         HttpResponse response = sendLoginPostForOtp(client, sessionDataKey, mockSMSProvider.getOTP());
+        EntityUtils.consume(response.getEntity());
 
-        authorizationCode = EntityUtils.toString(response.getEntity());
+        authorizationCode = mockClientCallback.getAuthorizationCode();
         assertNotNull(authorizationCode);
     }
 

diff --git a/...sts-backend/src/test/java/org/wso2/identity/integration/test/base/MockClientCallback.java b/...sts-backend/src/test/java/org/wso2/identity/integration/test/base/MockClientCallback.java
@@ -20,7 +20,10 @@
 
 import com.github.tomakehurst.wiremock.WireMockServer;
 import com.github.tomakehurst.wiremock.core.WireMockConfiguration;
+import com.github.tomakehurst.wiremock.extension.ResponseTransformerV2;
 import com.github.tomakehurst.wiremock.extension.responsetemplating.ResponseTemplateTransformer;
+import com.github.tomakehurst.wiremock.http.Response;
+import com.github.tomakehurst.wiremock.stubbing.ServeEvent;
 import org.wso2.identity.integration.common.utils.ISIntegrationTest;
 import org.wso2.identity.integration.test.util.Utils;
 
@@ -39,6 +42,8 @@ public class MockClientCallback {
 
     public static final String CALLBACK_URL = "https://localhost:8091/dummyApp/oauth2client";
 
+    private final AtomicReference<String> authorizationCode = new AtomicReference<>();
+
     private WireMockServer wireMockServer;
 
     public void start() {
@@ -50,7 +55,26 @@ public void start() {
                         ISIntegrationTest.KEYSTORE_NAME).toAbsolutePath().toString())
                 .keystorePassword("wso2carbon")
                 .keyManagerPassword("wso2carbon")
-                .extensions(new ResponseTemplateTransformer(null, true, null, null)));
+                .extensions(new ResponseTemplateTransformer(null, true, null, null),
+                        new ResponseTransformerV2() {
+
+                            @Override
+                            public Response transform(Response response, ServeEvent serveEvent) {
+
+                                authorizationCode.set(serveEvent.getRequest().getQueryParams().get("code").firstValue());
+                                return response;
+                            }
+
+                            @Override
+                            public boolean applyGlobally() {
+                                return false;
+                            }
+
+                            @Override
+                            public String getName() {
+                                return "authz-code-transformer";
+                            }
+                        }));
 
         wireMockServer.start();
 
@@ -71,11 +95,15 @@ private void configureMockEndpoints() {
             wireMockServer.stubFor(get(urlPathEqualTo("/dummyApp/oauth2client"))
                     .withQueryParam("code", matching(".*"))
                     .willReturn(aResponse()
-                            .withBody("{{request.query.code}}")
-                            .withTransformers("response-template")
+                            .withTransformers("response-template, authz-code-transformer")
                             .withStatus(200)));
         } catch (Exception e) {
             throw new RuntimeException(e);
         }
     }
+
+    public String getAuthorizationCode() {
+
+        return authorizationCode.get();
+    }
 }