PM-11155: Add logic for handling remove password flow (#3801)

2024-11-24 10:25:57 +03:00 · 2024-08-21 14:30:27 -05:00 · 2024-08-21 14:30:27 -05:00 · 13b256d4e9
commit 13b256d4e9
parent 5761e9510a
9 changed files with 336 additions and 4 deletions
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/repository/AuthRepository.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/repository/AuthRepository.kt
@ -16,6 +16,7 @@ import com.x8bit.bitwarden.data.auth.repository.model.PasswordStrengthResult
 import com.x8bit.bitwarden.data.auth.repository.model.PolicyInformation
 import com.x8bit.bitwarden.data.auth.repository.model.PrevalidateSsoResult
 import com.x8bit.bitwarden.data.auth.repository.model.RegisterResult
+import com.x8bit.bitwarden.data.auth.repository.model.RemovePasswordResult
 import com.x8bit.bitwarden.data.auth.repository.model.RequestOtpResult
 import com.x8bit.bitwarden.data.auth.repository.model.ResendEmailResult
 import com.x8bit.bitwarden.data.auth.repository.model.ResetPasswordResult
@ -267,6 +268,12 @@ interface AuthRepository : AuthenticatorProvider, AuthRequestManager {
        email: String,
    ): PasswordHintResult

+    /**
+     * Removes the users password from the account. This used used when migrating from master
+     * password login to key connector login.
+     */
+    suspend fun removePassword(masterPassword: String): RemovePasswordResult
+
    /**
     * Resets the users password from the [currentPassword] (or null for account recovery resets),
     * to the [newPassword] and optional [passwordHint].
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryImpl.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryImpl.kt
@ -35,6 +35,7 @@ import com.x8bit.bitwarden.data.auth.datasource.sdk.AuthSdkSource
 import com.x8bit.bitwarden.data.auth.datasource.sdk.util.toInt
 import com.x8bit.bitwarden.data.auth.datasource.sdk.util.toKdfTypeJson
 import com.x8bit.bitwarden.data.auth.manager.AuthRequestManager
+import com.x8bit.bitwarden.data.auth.manager.KeyConnectorManager
 import com.x8bit.bitwarden.data.auth.manager.TrustedDeviceManager
 import com.x8bit.bitwarden.data.auth.manager.UserLogoutManager
 import com.x8bit.bitwarden.data.auth.repository.model.AuthState
@ -49,6 +50,7 @@ import com.x8bit.bitwarden.data.auth.repository.model.PasswordStrengthResult
 import com.x8bit.bitwarden.data.auth.repository.model.PolicyInformation
 import com.x8bit.bitwarden.data.auth.repository.model.PrevalidateSsoResult
 import com.x8bit.bitwarden.data.auth.repository.model.RegisterResult
+import com.x8bit.bitwarden.data.auth.repository.model.RemovePasswordResult
 import com.x8bit.bitwarden.data.auth.repository.model.RequestOtpResult
 import com.x8bit.bitwarden.data.auth.repository.model.ResendEmailResult
 import com.x8bit.bitwarden.data.auth.repository.model.ResetPasswordResult
@ -95,6 +97,7 @@ import com.x8bit.bitwarden.data.platform.repository.util.bufferedMutableSharedFl
 import com.x8bit.bitwarden.data.platform.util.asFailure
 import com.x8bit.bitwarden.data.platform.util.asSuccess
 import com.x8bit.bitwarden.data.platform.util.flatMap
+import com.x8bit.bitwarden.data.vault.datasource.network.model.OrganizationType
 import com.x8bit.bitwarden.data.vault.datasource.network.model.PolicyTypeJson
 import com.x8bit.bitwarden.data.vault.datasource.network.model.SyncResponseJson
 import com.x8bit.bitwarden.data.vault.datasource.sdk.VaultSdkSource
@ -145,6 +148,7 @@ class AuthRepositoryImpl(
    private val settingsRepository: SettingsRepository,
    private val vaultRepository: VaultRepository,
    private val authRequestManager: AuthRequestManager,
+    private val keyConnectorManager: KeyConnectorManager,
    private val trustedDeviceManager: TrustedDeviceManager,
    private val userLogoutManager: UserLogoutManager,
    private val policyManager: PolicyManager,
@ -841,6 +845,40 @@ class AuthRepositoryImpl(
        )
    }

+    override suspend fun removePassword(masterPassword: String): RemovePasswordResult {
+        val activeAccount = authDiskSource
+            .userState
+            ?.activeAccount
+            ?: return RemovePasswordResult.Error
+        val profile = activeAccount.profile
+        val userId = profile.userId
+        val userKey = authDiskSource
+            .getUserKey(userId = userId)
+            ?: return RemovePasswordResult.Error
+        val keyConnectorUrl = organizations
+            .find {
+                it.shouldUseKeyConnector &&
+                    it.type != OrganizationType.OWNER &&
+                    it.type != OrganizationType.ADMIN
+            }
+            ?.keyConnectorUrl
+            ?: return RemovePasswordResult.Error
+        return keyConnectorManager
+            .migrateExistingUserToKeyConnector(
+                userId = userId,
+                url = keyConnectorUrl,
+                userKeyEncrypted = userKey,
+                email = profile.email,
+                masterPassword = masterPassword,
+                kdf = profile.toSdkParams(),
+            )
+            .onSuccess { vaultRepository.sync() }
+            .fold(
+                onFailure = { RemovePasswordResult.Error },
+                onSuccess = { RemovePasswordResult.Success },
+            )
+    }
+
    override suspend fun resetPassword(
        currentPassword: String?,
        newPassword: String,
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/repository/di/AuthRepositoryModule.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/repository/di/AuthRepositoryModule.kt
@ -8,6 +8,7 @@ import com.x8bit.bitwarden.data.auth.datasource.network.service.IdentityService
 import com.x8bit.bitwarden.data.auth.datasource.network.service.OrganizationService
 import com.x8bit.bitwarden.data.auth.datasource.sdk.AuthSdkSource
 import com.x8bit.bitwarden.data.auth.manager.AuthRequestManager
+import com.x8bit.bitwarden.data.auth.manager.KeyConnectorManager
 import com.x8bit.bitwarden.data.auth.manager.TrustedDeviceManager
 import com.x8bit.bitwarden.data.auth.manager.UserLogoutManager
 import com.x8bit.bitwarden.data.auth.repository.AuthRepository
@ -48,6 +49,7 @@ object AuthRepositoryModule {
        environmentRepository: EnvironmentRepository,
        settingsRepository: SettingsRepository,
        vaultRepository: VaultRepository,
+        keyConnectorManager: KeyConnectorManager,
        authRequestManager: AuthRequestManager,
        trustedDeviceManager: TrustedDeviceManager,
        userLogoutManager: UserLogoutManager,
@ -67,6 +69,7 @@ object AuthRepositoryModule {
        environmentRepository = environmentRepository,
        settingsRepository = settingsRepository,
        vaultRepository = vaultRepository,
+        keyConnectorManager = keyConnectorManager,
        authRequestManager = authRequestManager,
        trustedDeviceManager = trustedDeviceManager,
        userLogoutManager = userLogoutManager,
--- a/app/src/main/java/com/x8bit/bitwarden/data/auth/repository/model/RemovePasswordResult.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/data/auth/repository/model/RemovePasswordResult.kt
@ -0,0 +1,16 @@
+package com.x8bit.bitwarden.data.auth.repository.model
+
+/**
+ * Models result of removing a user's password.
+ */
+sealed class RemovePasswordResult {
+    /**
+     * The password was removed successfully.
+     */
+    data object Success : RemovePasswordResult()
+
+    /**
+     * There was an error removing the password.
+     */
+    data object Error : RemovePasswordResult()
+}
--- a/app/src/main/java/com/x8bit/bitwarden/ui/auth/feature/removepassword/RemovePasswordScreen.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/ui/auth/feature/removepassword/RemovePasswordScreen.kt
@ -32,6 +32,8 @@ import com.x8bit.bitwarden.ui.platform.components.appbar.BitwardenTopAppBar
 import com.x8bit.bitwarden.ui.platform.components.button.BitwardenFilledButton
 import com.x8bit.bitwarden.ui.platform.components.dialog.BasicDialogState
 import com.x8bit.bitwarden.ui.platform.components.dialog.BitwardenBasicDialog
+import com.x8bit.bitwarden.ui.platform.components.dialog.BitwardenLoadingDialog
+import com.x8bit.bitwarden.ui.platform.components.dialog.LoadingDialogState
 import com.x8bit.bitwarden.ui.platform.components.field.BitwardenPasswordField
 import com.x8bit.bitwarden.ui.platform.components.scaffold.BitwardenScaffold
 import com.x8bit.bitwarden.ui.platform.theme.BitwardenTheme
@ -147,6 +149,12 @@ private fun RemovePasswordDialogs(
            )
        }

+        is RemovePasswordState.DialogState.Loading -> {
+            BitwardenLoadingDialog(
+                visibilityState = LoadingDialogState.Shown(text = dialogState.title),
+            )
+        }
+
        null -> Unit
    }
 }
--- a/app/src/main/java/com/x8bit/bitwarden/ui/auth/feature/removepassword/RemovePasswordViewModel.kt
+++ b/app/src/main/java/com/x8bit/bitwarden/ui/auth/feature/removepassword/RemovePasswordViewModel.kt
@ -2,13 +2,16 @@ package com.x8bit.bitwarden.ui.auth.feature.removepassword

 import android.os.Parcelable
 import androidx.lifecycle.SavedStateHandle
+import androidx.lifecycle.viewModelScope
 import com.x8bit.bitwarden.R
 import com.x8bit.bitwarden.data.auth.repository.AuthRepository
+import com.x8bit.bitwarden.data.auth.repository.model.RemovePasswordResult
 import com.x8bit.bitwarden.ui.platform.base.BaseViewModel
 import com.x8bit.bitwarden.ui.platform.base.util.Text
 import com.x8bit.bitwarden.ui.platform.base.util.asText
 import dagger.hilt.android.lifecycle.HiltViewModel
 import kotlinx.coroutines.flow.update
+import kotlinx.coroutines.launch
 import kotlinx.parcelize.Parcelize
 import javax.inject.Inject

@ -43,11 +46,36 @@ class RemovePasswordViewModel @Inject constructor(
            RemovePasswordAction.ContinueClick -> handleContinueClick()
            is RemovePasswordAction.InputChanged -> handleInputChanged(action)
            RemovePasswordAction.DialogDismiss -> handleDialogDismiss()
+            is RemovePasswordAction.Internal.ReceiveRemovePasswordResult -> {
+                handleReceiveRemovePasswordResult(action)
+            }
        }
    }

    private fun handleContinueClick() {
-        // TODO: Process removing the password (PM-11155)
+        if (state.input.isBlank()) {
+            mutableStateFlow.update {
+                it.copy(
+                    dialogState = RemovePasswordState.DialogState.Error(
+                        title = R.string.an_error_has_occurred.asText(),
+                        message = R.string.validation_field_required
+                            .asText(R.string.master_password.asText()),
+                    ),
+                )
+            }
+            return
+        }
+        mutableStateFlow.update {
+            it.copy(
+                dialogState = RemovePasswordState.DialogState.Loading(
+                    title = R.string.deleting.asText(),
+                ),
+            )
+        }
+        viewModelScope.launch {
+            val result = authRepository.removePassword(masterPassword = state.input)
+            sendAction(RemovePasswordAction.Internal.ReceiveRemovePasswordResult(result))
+        }
    }

    private fun handleInputChanged(action: RemovePasswordAction.InputChanged) {
@ -57,6 +85,28 @@ class RemovePasswordViewModel @Inject constructor(
    private fun handleDialogDismiss() {
        mutableStateFlow.update { it.copy(dialogState = null) }
    }
+
+    private fun handleReceiveRemovePasswordResult(
+        action: RemovePasswordAction.Internal.ReceiveRemovePasswordResult,
+    ) {
+        when (action.result) {
+            RemovePasswordResult.Error -> {
+                mutableStateFlow.update {
+                    it.copy(
+                        dialogState = RemovePasswordState.DialogState.Error(
+                            title = R.string.an_error_has_occurred.asText(),
+                            message = R.string.generic_error_message.asText(),
+                        ),
+                    )
+                }
+            }
+
+            RemovePasswordResult.Success -> {
+                mutableStateFlow.update { it.copy(dialogState = null) }
+                // We do nothing here because state-based navigation will handle it.
+            }
+        }
+    }
 }

 /**
@ -81,6 +131,12 @@ data class RemovePasswordState(
            val title: Text? = null,
            val message: Text,
        ) : DialogState()
+
+        /**
+         * Represents a loading dialog with the given [title].
+         */
+        @Parcelize
+        data class Loading(val title: Text) : DialogState()
    }
 }

@ -104,4 +160,16 @@ sealed class RemovePasswordAction {
     * Indicates that the dialog has been dismissed.
     */
    data object DialogDismiss : RemovePasswordAction()
+
+    /**
+     * Models actions that the [RemovePasswordViewModel] might send itself.
+     */
+    sealed class Internal : RemovePasswordAction() {
+        /**
+         * Indicates that a remove password result has been received.
+         */
+        data class ReceiveRemovePasswordResult(
+            val result: RemovePasswordResult,
+        ) : Internal()
+    }
 }
--- a/app/src/test/java/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryTest.kt
+++ b/app/src/test/java/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryTest.kt
@ -52,6 +52,7 @@ import com.x8bit.bitwarden.data.auth.datasource.sdk.model.PasswordStrength.LEVEL
 import com.x8bit.bitwarden.data.auth.datasource.sdk.model.PasswordStrength.LEVEL_3
 import com.x8bit.bitwarden.data.auth.datasource.sdk.model.PasswordStrength.LEVEL_4
 import com.x8bit.bitwarden.data.auth.manager.AuthRequestManager
+import com.x8bit.bitwarden.data.auth.manager.KeyConnectorManager
 import com.x8bit.bitwarden.data.auth.manager.TrustedDeviceManager
 import com.x8bit.bitwarden.data.auth.manager.UserLogoutManager
 import com.x8bit.bitwarden.data.auth.manager.model.AuthRequest
@ -66,6 +67,7 @@ import com.x8bit.bitwarden.data.auth.repository.model.PasswordHintResult
 import com.x8bit.bitwarden.data.auth.repository.model.PasswordStrengthResult
 import com.x8bit.bitwarden.data.auth.repository.model.PrevalidateSsoResult
 import com.x8bit.bitwarden.data.auth.repository.model.RegisterResult
+import com.x8bit.bitwarden.data.auth.repository.model.RemovePasswordResult
 import com.x8bit.bitwarden.data.auth.repository.model.RequestOtpResult
 import com.x8bit.bitwarden.data.auth.repository.model.ResendEmailResult
 import com.x8bit.bitwarden.data.auth.repository.model.ResetPasswordResult
@ -100,6 +102,7 @@ import com.x8bit.bitwarden.data.platform.repository.util.FakeEnvironmentReposito
 import com.x8bit.bitwarden.data.platform.repository.util.bufferedMutableSharedFlow
 import com.x8bit.bitwarden.data.platform.util.asFailure
 import com.x8bit.bitwarden.data.platform.util.asSuccess
+import com.x8bit.bitwarden.data.vault.datasource.network.model.OrganizationType
 import com.x8bit.bitwarden.data.vault.datasource.network.model.PolicyTypeJson
 import com.x8bit.bitwarden.data.vault.datasource.network.model.SyncResponseJson
 import com.x8bit.bitwarden.data.vault.datasource.network.model.createMockOrganization
@ -206,6 +209,7 @@ class AuthRepositoryTest {
        } returns "AsymmetricEncString".asSuccess()
    }
    private val authRequestManager: AuthRequestManager = mockk()
+    private val keyConnectorManager: KeyConnectorManager = mockk()
    private val trustedDeviceManager: TrustedDeviceManager = mockk()
    private val userLogoutManager: UserLogoutManager = mockk {
        every { logout(any(), any()) } just runs
@ -238,6 +242,7 @@ class AuthRepositoryTest {
        settingsRepository = settingsRepository,
        vaultRepository = vaultRepository,
        authRequestManager = authRequestManager,
+        keyConnectorManager = keyConnectorManager,
        trustedDeviceManager = trustedDeviceManager,
        userLogoutManager = userLogoutManager,
        dispatcherManager = dispatcherManager,
@ -3810,6 +3815,114 @@ class AuthRepositoryTest {
        assertEquals(RegisterResult.Success(CAPTCHA_KEY), result)
    }

+    @Test
+    fun `removePassword with no active account should return error`() = runTest {
+        fakeAuthDiskSource.userState = null
+
+        val result = repository.removePassword(masterPassword = PASSWORD)
+
+        assertEquals(RemovePasswordResult.Error, result)
+    }
+
+    @Test
+    fun `removePassword with no userKey should return error`() = runTest {
+        fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
+        fakeAuthDiskSource.storeUserKey(userId = USER_ID_1, userKey = null)
+
+        val result = repository.removePassword(masterPassword = PASSWORD)
+
+        assertEquals(RemovePasswordResult.Error, result)
+    }
+
+    @Test
+    fun `removePassword with no keyConnectorUrl should return error`() = runTest {
+        fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
+        fakeAuthDiskSource.storeUserKey(userId = USER_ID_1, userKey = ENCRYPTED_USER_KEY)
+        val organizations = listOf(
+            mockk<SyncResponseJson.Profile.Organization> {
+                every { id } returns "orgId"
+                every { name } returns "orgName"
+                every { shouldUseKeyConnector } returns true
+                every { type } returns OrganizationType.USER
+                every { keyConnectorUrl } returns null
+            },
+        )
+        fakeAuthDiskSource.storeOrganizations(userId = USER_ID_1, organizations = organizations)
+
+        val result = repository.removePassword(masterPassword = PASSWORD)
+
+        assertEquals(RemovePasswordResult.Error, result)
+    }
+
+    @Test
+    fun `removePassword with migrateExistingUserToKeyConnector error should return error`() =
+        runTest {
+            fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
+            fakeAuthDiskSource.storeUserKey(userId = USER_ID_1, userKey = ENCRYPTED_USER_KEY)
+            val url = "www.example.com"
+            val organizations = listOf(
+                mockk<SyncResponseJson.Profile.Organization> {
+                    every { id } returns "orgId"
+                    every { name } returns "orgName"
+                    every { shouldUseKeyConnector } returns true
+                    every { type } returns OrganizationType.USER
+                    every { keyConnectorUrl } returns url
+                },
+            )
+            fakeAuthDiskSource.storeOrganizations(userId = USER_ID_1, organizations = organizations)
+            coEvery {
+                keyConnectorManager.migrateExistingUserToKeyConnector(
+                    userId = USER_ID_1,
+                    url = url,
+                    userKeyEncrypted = ENCRYPTED_USER_KEY,
+                    email = PROFILE_1.email,
+                    masterPassword = PASSWORD,
+                    kdf = PROFILE_1.toSdkParams(),
+                )
+            } returns Throwable("Fail").asFailure()
+
+            val result = repository.removePassword(masterPassword = PASSWORD)
+
+            assertEquals(RemovePasswordResult.Error, result)
+        }
+
+    @Suppress("MaxLineLength")
+    @Test
+    fun `removePassword with migrateExistingUserToKeyConnector success should sync and return success`() =
+        runTest {
+            fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
+            fakeAuthDiskSource.storeUserKey(userId = USER_ID_1, userKey = ENCRYPTED_USER_KEY)
+            val url = "www.example.com"
+            val organizations = listOf(
+                mockk<SyncResponseJson.Profile.Organization> {
+                    every { id } returns "orgId"
+                    every { name } returns "orgName"
+                    every { shouldUseKeyConnector } returns true
+                    every { type } returns OrganizationType.USER
+                    every { keyConnectorUrl } returns url
+                },
+            )
+            fakeAuthDiskSource.storeOrganizations(userId = USER_ID_1, organizations = organizations)
+            coEvery {
+                keyConnectorManager.migrateExistingUserToKeyConnector(
+                    userId = USER_ID_1,
+                    url = url,
+                    userKeyEncrypted = ENCRYPTED_USER_KEY,
+                    email = PROFILE_1.email,
+                    masterPassword = PASSWORD,
+                    kdf = PROFILE_1.toSdkParams(),
+                )
+            } returns Unit.asSuccess()
+            every { vaultRepository.sync() } just runs
+
+            val result = repository.removePassword(masterPassword = PASSWORD)
+
+            assertEquals(RemovePasswordResult.Success, result)
+            verify(exactly = 1) {
+                vaultRepository.sync()
+            }
+        }
+
    @Test
    fun `resetPassword Success should return Success`() = runTest {
        val currentPassword = "currentPassword"
--- a/app/src/test/java/com/x8bit/bitwarden/ui/auth/feature/removepassword/RemovePasswordScreenTest.kt
+++ b/app/src/test/java/com/x8bit/bitwarden/ui/auth/feature/removepassword/RemovePasswordScreenTest.kt
@ -63,6 +63,20 @@ class RemovePasswordScreenTest : BaseComposeTest() {
            .assert(hasAnyAncestor(isDialog()))
            .isDisplayed()

+        val loadingMessage = "Loading message"
+        mutableStateFlow.update {
+            it.copy(
+                dialogState = RemovePasswordState.DialogState.Loading(
+                    title = loadingMessage.asText(),
+                ),
+            )
+        }
+
+        composeTestRule
+            .onNodeWithText(text = loadingMessage)
+            .assert(hasAnyAncestor(isDialog()))
+            .isDisplayed()
+
        mutableStateFlow.update { it.copy(dialogState = null) }

        composeTestRule.onNode(isDialog()).assertDoesNotExist()
--- a/app/src/test/java/com/x8bit/bitwarden/ui/auth/feature/removepassword/RemovePasswordViewModelTest.kt
+++ b/app/src/test/java/com/x8bit/bitwarden/ui/auth/feature/removepassword/RemovePasswordViewModelTest.kt
@ -5,11 +5,13 @@ import app.cash.turbine.test
 import com.x8bit.bitwarden.R
 import com.x8bit.bitwarden.data.auth.repository.AuthRepository
 import com.x8bit.bitwarden.data.auth.repository.model.Organization
+import com.x8bit.bitwarden.data.auth.repository.model.RemovePasswordResult
 import com.x8bit.bitwarden.data.auth.repository.model.UserState
 import com.x8bit.bitwarden.data.platform.repository.model.Environment
 import com.x8bit.bitwarden.data.vault.datasource.network.model.OrganizationType
 import com.x8bit.bitwarden.ui.platform.base.BaseViewModelTest
 import com.x8bit.bitwarden.ui.platform.base.util.asText
+import io.mockk.coEvery
 import io.mockk.every
 import io.mockk.mockk
 import kotlinx.coroutines.flow.MutableStateFlow
@ -24,11 +26,74 @@ class RemovePasswordViewModelTest : BaseViewModelTest() {
    }

    @Test
-    fun `ContinueClick calls does nothing`() = runTest {
+    fun `ContinueClick with blank input should show error dialog`() {
        val viewModel = createViewModel()
-        viewModel.eventFlow.test {
        viewModel.trySendAction(RemovePasswordAction.ContinueClick)
-            expectNoEvents()
+        assertEquals(
+            DEFAULT_STATE.copy(
+                dialogState = RemovePasswordState.DialogState.Error(
+                    title = R.string.an_error_has_occurred.asText(),
+                    message = R.string.validation_field_required
+                        .asText(R.string.master_password.asText()),
+                ),
+            ),
+            viewModel.stateFlow.value,
+        )
+    }
+
+    @Test
+    fun `ContinueClick with input and remove password error should show error dialog`() = runTest {
+        val password = "123"
+        val initialState = DEFAULT_STATE.copy(input = password)
+        val viewModel = createViewModel(state = initialState)
+        coEvery {
+            authRepository.removePassword(masterPassword = password)
+        } returns RemovePasswordResult.Error
+
+        viewModel.stateFlow.test {
+            assertEquals(initialState, awaitItem())
+            viewModel.trySendAction(RemovePasswordAction.ContinueClick)
+            assertEquals(
+                initialState.copy(
+                    dialogState = RemovePasswordState.DialogState.Loading(
+                        title = R.string.deleting.asText(),
+                    ),
+                ),
+                awaitItem(),
+            )
+            assertEquals(
+                initialState.copy(
+                    dialogState = RemovePasswordState.DialogState.Error(
+                        title = R.string.an_error_has_occurred.asText(),
+                        message = R.string.generic_error_message.asText(),
+                    ),
+                ),
+                awaitItem(),
+            )
+        }
+    }
+
+    @Test
+    fun `ContinueClick with input and remove password success should dismiss dialog`() = runTest {
+        val password = "123"
+        val initialState = DEFAULT_STATE.copy(input = password)
+        val viewModel = createViewModel(state = initialState)
+        coEvery {
+            authRepository.removePassword(masterPassword = password)
+        } returns RemovePasswordResult.Success
+
+        viewModel.stateFlow.test {
+            assertEquals(initialState, awaitItem())
+            viewModel.trySendAction(RemovePasswordAction.ContinueClick)
+            assertEquals(
+                initialState.copy(
+                    dialogState = RemovePasswordState.DialogState.Loading(
+                        title = R.string.deleting.asText(),
+                    ),
+                ),
+                awaitItem(),
+            )
+            assertEquals(initialState, awaitItem())
        }
    }