fix(auth): Apply Okta migration workaround to 2FA check (#30276)

src/sentry/auth/helper.py

@@ -80,6 +80,16 @@ class AuthHelperSessionStore(PipelineSessionStore):
         self.request.session.modified = True
 
 
+def _using_okta_migration_workaround(
+    organization: Organization, user: User, auth_provider: Optional[AuthProvider]
+) -> bool:
+    # XXX(leedongwei): Workaround for migrating Okta instance
+    # TODO: Delete after workaround is not needed
+    has_flag = features.has("organizations:sso-migration", organization, actor=user)
+    has_provider = auth_provider and (auth_provider.provider in ("okta", "saml2"))
+    return has_flag and has_provider
+
+
 Identity = Mapping[str, Any]
 
 
@@ -440,8 +450,16 @@ class AuthIdentityHandler:
                 try:
                     self._login(acting_user)
                 except self._NotCompletedSecurityChecks:
-                    # adding is_account_verified to the check below in order to redirect to 2fa when the user migrates their idp but has 2fa enabled, otherwise it would stop them from linking their sso provider
-                    if acting_user.has_usable_password() or is_account_verified:
+                    # adding is_account_verified to the check below in order to redirect
+                    # to 2fa when the user migrates their idp but has 2fa enabled,
+                    # otherwise it would stop them from linking their sso provider
+                    if (
+                        acting_user.has_usable_password()
+                        or is_account_verified
+                        or _using_okta_migration_workaround(
+                            self.organization, acting_user, self.auth_provider
+                        )
+                    ):
                         return self._post_login_redirect()
                     else:
                         acting_user = None
@@ -455,6 +473,7 @@ class AuthIdentityHandler:
         elif acting_user and not acting_user.has_usable_password():
             acting_user = None
 
+        auth_identity = None
         if op == "confirm" and self.user.is_authenticated:
             auth_identity = self.handle_attach_identity(identity)
         elif op == "newuser":
@@ -674,6 +693,8 @@ class AuthHelper(Pipeline):
             response = self._finish_login_pipeline(identity)
         elif self.state.flow == self.FLOW_SETUP_PROVIDER:
             response = self._finish_setup_pipeline(identity)
+        else:
+            raise Exception(f"Unrecognized flow value: {self.state.flow}")
 
         return response
 
@@ -721,10 +742,9 @@ class AuthHelper(Pipeline):
                     auth_identity = None
 
             if not auth_identity:
-                # XXX(leedongwei): Workaround for migrating Okta instance
-                if features.has(
-                    "organizations:sso-migration", self.organization, actor=self.request.user
-                ) and (auth_provider.provider == "okta" or auth_provider.provider == "saml2"):
+                if _using_okta_migration_workaround(
+                    self.organization, self.request.user, auth_provider
+                ):
                     identity["email_verified"] = True
 
                     logger.info(