From 721c69619e7ee676dece52ae3585ac155eff1784 Mon Sep 17 00:00:00 2001
From: Dave Severns <149429124+dseverns-livefront@users.noreply.github.com>
Date: Mon, 15 Jul 2024 14:32:40 -0400
Subject: [PATCH] PM-9017 updated the continue button state when switching 2FA
 method (#3530)

---
 .../twofactorlogin/TwoFactorLoginViewModel.kt | 13 +++++++++--
 .../TwoFactorLoginViewModelTest.kt            | 23 +++++++++++++++++--
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/app/src/main/java/com/x8bit/bitwarden/ui/auth/feature/twofactorlogin/TwoFactorLoginViewModel.kt b/app/src/main/java/com/x8bit/bitwarden/ui/auth/feature/twofactorlogin/TwoFactorLoginViewModel.kt
index 5c569bab3..53639b9fd 100644
--- a/app/src/main/java/com/x8bit/bitwarden/ui/auth/feature/twofactorlogin/TwoFactorLoginViewModel.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/ui/auth/feature/twofactorlogin/TwoFactorLoginViewModel.kt
@@ -465,7 +465,7 @@ class TwoFactorLoginViewModel @Inject constructor(
                         )
                     }
                 }
-                mutableStateFlow.update { it.copy(authMethod = action.authMethod) }
+                updateAuthMethodRelatedState(action.authMethod)
             }
 
             TwoFactorAuthMethod.AUTHENTICATOR_APP,
@@ -476,11 +476,20 @@ class TwoFactorLoginViewModel @Inject constructor(
             TwoFactorAuthMethod.DUO_ORGANIZATION,
             TwoFactorAuthMethod.WEB_AUTH,
             -> {
-                mutableStateFlow.update { it.copy(authMethod = action.authMethod) }
+                updateAuthMethodRelatedState(action.authMethod)
             }
         }
     }
 
+    private fun updateAuthMethodRelatedState(authMethod: TwoFactorAuthMethod) {
+        mutableStateFlow.update {
+            it.copy(
+                authMethod = authMethod,
+                isContinueButtonEnabled = authMethod.isContinueButtonEnabled,
+            )
+        }
+    }
+
     /**
      * Verify the input and attempt to authenticate with the code.
      */
diff --git a/app/src/test/java/com/x8bit/bitwarden/ui/auth/feature/twofactorlogin/TwoFactorLoginViewModelTest.kt b/app/src/test/java/com/x8bit/bitwarden/ui/auth/feature/twofactorlogin/TwoFactorLoginViewModelTest.kt
index 455f0b60e..b6e72d92b 100644
--- a/app/src/test/java/com/x8bit/bitwarden/ui/auth/feature/twofactorlogin/TwoFactorLoginViewModelTest.kt
+++ b/app/src/test/java/com/x8bit/bitwarden/ui/auth/feature/twofactorlogin/TwoFactorLoginViewModelTest.kt
@@ -680,14 +680,33 @@ class TwoFactorLoginViewModelTest : BaseViewModelTest() {
     @Test
     fun `SelectAuthMethod with other method should update the state`() {
         val viewModel = createViewModel()
+
+        // To method with continue button enabled by default
         viewModel.trySendAction(
             TwoFactorLoginAction.SelectAuthMethod(
-                TwoFactorAuthMethod.AUTHENTICATOR_APP,
+                TwoFactorAuthMethod.DUO,
             ),
         )
+
         assertEquals(
             DEFAULT_STATE.copy(
-                authMethod = TwoFactorAuthMethod.AUTHENTICATOR_APP,
+                authMethod = TwoFactorAuthMethod.DUO,
+                isContinueButtonEnabled = true,
+            ),
+            viewModel.stateFlow.value,
+        )
+
+        // To method with continue button disabled by default
+        viewModel.trySendAction(
+            TwoFactorLoginAction.SelectAuthMethod(
+                TwoFactorAuthMethod.YUBI_KEY,
+            ),
+        )
+
+        assertEquals(
+            DEFAULT_STATE.copy(
+                authMethod = TwoFactorAuthMethod.YUBI_KEY,
+                isContinueButtonEnabled = false,
             ),
             viewModel.stateFlow.value,
         )