From 197df340973f8d0490791e6c23e6d077ffbc5406 Mon Sep 17 00:00:00 2001
From: Adam Brown <adampsbrown@gmail.com>
Date: Thu, 14 Apr 2022 10:03:45 +0100
Subject: [PATCH] only setting selected homeserver state after a successful
 start registration when editing

---
 .../onboarding/OnboardingViewModel.kt         | 25 +++++++++++++------
 1 file changed, 17 insertions(+), 8 deletions(-)

diff --git a/vector/src/main/java/im/vector/app/features/onboarding/OnboardingViewModel.kt b/vector/src/main/java/im/vector/app/features/onboarding/OnboardingViewModel.kt
index 7ce25190cd..448fd514ce 100644
--- a/vector/src/main/java/im/vector/app/features/onboarding/OnboardingViewModel.kt
+++ b/vector/src/main/java/im/vector/app/features/onboarding/OnboardingViewModel.kt
@@ -628,24 +628,33 @@ class OnboardingViewModel @AssistedInject constructor(
             _viewEvents.post(OnboardingViewEvents.OutdatedHomeserver)
         }
 
-        setState {
-            copy(
-                    serverType = alignServerTypeAfterSubmission(config, serverTypeOverride),
-                    selectedHomeserver = authResult.selectedHomeserver,
-            )
-        }
         val state = awaitState()
+
         when (lastAction) {
             is OnboardingAction.HomeServerChange.EditHomeServer   -> {
                 when (state.onboardingFlow) {
-                    OnboardingFlow.SignUp -> handleRegisterAction(RegisterAction.StartRegistration) { _ ->
+                    OnboardingFlow.SignUp -> internalRegisterAction(RegisterAction.StartRegistration) { _ ->
+                        setState {
+                            copy(
+                                    serverType = alignServerTypeAfterSubmission(config, serverTypeOverride),
+                                    selectedHomeserver = authResult.selectedHomeserver,
+                            )
+                        }
                         _viewEvents.post(OnboardingViewEvents.OnHomeserverEdited)
                     }
                     else                  -> throw IllegalArgumentException("developer error")
                 }
             }
             is OnboardingAction.HomeServerChange.SelectHomeServer -> {
-                if (state.selectedHomeserver.preferredLoginMode.supportsSignModeScreen()) {
+                setState {
+                    copy(
+                            serverType = alignServerTypeAfterSubmission(config, serverTypeOverride),
+                            selectedHomeserver = authResult.selectedHomeserver,
+                    )
+                }
+
+
+                if (authResult.selectedHomeserver.preferredLoginMode.supportsSignModeScreen()) {
                     when (state.onboardingFlow) {
                         OnboardingFlow.SignIn -> internalRegisterAction(RegisterAction.StartRegistration, ::emitFlowResultViewEvent)
                         OnboardingFlow.SignUp -> internalRegisterAction(RegisterAction.StartRegistration, ::emitFlowResultViewEvent)