diff --git a/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/auth/StageExtension.kt b/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/auth/StageExtension.kt
new file mode 100644
index 0000000000000000000000000000000000000000..4f34a37fb6b37a86d7495225d3d84670e9039a6b
--- /dev/null
+++ b/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/auth/StageExtension.kt
@@ -0,0 +1,14 @@
+package org.matrix.android.sdk.internal.auth
+
+import org.matrix.android.sdk.api.auth.data.LoginFlowTypes
+import org.matrix.android.sdk.api.auth.registration.Stage
+
+//Added to support few registration flows
+fun List<Stage>.findStageForType(type: String): Stage? = when (type) {
+    LoginFlowTypes.RECAPTCHA      -> firstOrNull { it is Stage.ReCaptcha }
+    LoginFlowTypes.DUMMY          -> firstOrNull { it is Stage.Dummy }
+    LoginFlowTypes.TERMS          -> firstOrNull { it is Stage.Terms }
+    LoginFlowTypes.EMAIL_IDENTITY -> firstOrNull { it is Stage.Email }
+    LoginFlowTypes.MSISDN         -> firstOrNull { it is Stage.Msisdn }
+    else                          -> firstOrNull { (it as? Stage.Other)?.type == type }
+}
\ No newline at end of file
diff --git a/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/auth/registration/DefaultRegistrationWizard.kt b/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/auth/registration/DefaultRegistrationWizard.kt
index 63bf1b0a260819c0767ade9f5fb908e79ee5b841..8401db7fbb8d014ebf019952265ed227e38696ea 100644
--- a/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/auth/registration/DefaultRegistrationWizard.kt
+++ b/matrix-sdk-android/src/main/java/org/matrix/android/sdk/internal/auth/registration/DefaultRegistrationWizard.kt
@@ -33,6 +33,7 @@ import org.matrix.android.sdk.internal.auth.AuthAPI
 import org.matrix.android.sdk.internal.auth.PendingSessionStore
 import org.matrix.android.sdk.internal.auth.SessionCreator
 import org.matrix.android.sdk.internal.auth.db.PendingSessionData
+import org.matrix.android.sdk.internal.auth.findStageForType
 
 /**
  * This class execute the registration request and is responsible to keep the session of interactive authentication.
@@ -295,16 +296,6 @@ internal class DefaultRegistrationWizard(
         return emptyList()
     }
 
-    //Added to support few registration flows
-    private fun List<Stage>.findStageForType(type: String): Stage? = when (type) {
-        LoginFlowTypes.RECAPTCHA      -> firstOrNull { it is Stage.ReCaptcha }
-        LoginFlowTypes.DUMMY          -> firstOrNull { it is Stage.Dummy }
-        LoginFlowTypes.TERMS          -> firstOrNull { it is Stage.Terms }
-        LoginFlowTypes.EMAIL_IDENTITY -> firstOrNull { it is Stage.Email }
-        LoginFlowTypes.MSISDN         -> firstOrNull { it is Stage.Msisdn }
-        else                          -> firstOrNull { (it as? Stage.Other)?.type == type }
-    }
-
     override suspend fun registrationSwiclops(
             authParams: JsonDict,
             userName: String?,