Add ability to set a PIN to SettingsRepository (#631)

app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/disk/AuthDiskSource.kt

@@ -94,8 +94,15 @@ interface AuthDiskSource {
 
     /**
      * Stores a pin-protected user key for the given [userId].
+     *
+     * When [inMemoryOnly] is `true`, the value will only be available via a call to
+     * [getPinProtectedUserKey] during the current app session.
      */
-    fun storePinProtectedUserKey(userId: String, pinProtectedUserKey: String?)
+    fun storePinProtectedUserKey(
+        userId: String,
+        pinProtectedUserKey: String?,
+        inMemoryOnly: Boolean = false,
+    )
 
     /**
      * Retrieves an encrypted PIN for the given [userId].

app/src/main/java/com/x8bit/bitwarden/data/auth/datasource/disk/AuthDiskSourceImpl.kt

@@ -40,6 +40,7 @@ class AuthDiskSourceImpl(
     sharedPreferences = sharedPreferences,
 ),
     AuthDiskSource {
+    private val inMemoryPinProtectedUserKeys = mutableMapOf<String, String?>()
     private val mutableOrganizationsFlowMap =
         mutableMapOf<String, MutableSharedFlow<List<SyncResponseJson.Profile.Organization>?>>()
 
@@ -132,12 +133,16 @@ class AuthDiskSourceImpl(
     }
 
     override fun getPinProtectedUserKey(userId: String): String? =
-        getString(key = "${PIN_PROTECTED_USER_KEY_KEY}_$userId")
+        inMemoryPinProtectedUserKeys[userId]
+            ?: getString(key = "${PIN_PROTECTED_USER_KEY_KEY}_$userId")
 
     override fun storePinProtectedUserKey(
         userId: String,
         pinProtectedUserKey: String?,
+        inMemoryOnly: Boolean,
     ) {
+        inMemoryPinProtectedUserKeys[userId] = pinProtectedUserKey
+        if (inMemoryOnly) return
         putString(
             key = "${PIN_PROTECTED_USER_KEY_KEY}_$userId",
             value = pinProtectedUserKey,

app/src/main/java/com/x8bit/bitwarden/data/platform/repository/SettingsRepository.kt

@@ -36,6 +36,11 @@ interface SettingsRepository {
      */
     var vaultTimeoutAction: VaultTimeoutAction
 
+    /**
+     * Whether or not PIN unlocking is enabled for the current user.
+     */
+    val isUnlockWithPinEnabled: Boolean
+
     /**
      * Clears all the settings data for the given user.
      */
@@ -85,4 +90,21 @@ interface SettingsRepository {
      * Stores the given [isPullToRefreshEnabled] for the active user.
      */
     fun storePullToRefreshEnabled(isPullToRefreshEnabled: Boolean)
+
+    /**
+     * Stores the given PIN, allowing it to be used to unlock the current user's vault.
+     *
+     * When [shouldRequireMasterPasswordOnRestart] is `true`, the user's master password is required
+     * on app startup but they may use their PIN to unlock their vault if it becomes locked while
+     * the app is still open.
+     */
+    fun storeUnlockPin(
+        pin: String,
+        shouldRequireMasterPasswordOnRestart: Boolean,
+    )
+
+    /**
+     * Clears any previously set unlock PIN for the current user.
+     */
+    fun clearUnlockPin()
 }

app/src/main/java/com/x8bit/bitwarden/data/platform/repository/SettingsRepositoryImpl.kt

@@ -5,6 +5,7 @@ import com.x8bit.bitwarden.data.platform.datasource.disk.SettingsDiskSource
 import com.x8bit.bitwarden.data.platform.manager.dispatcher.DispatcherManager
 import com.x8bit.bitwarden.data.platform.repository.model.VaultTimeout
 import com.x8bit.bitwarden.data.platform.repository.model.VaultTimeoutAction
+import com.x8bit.bitwarden.data.vault.datasource.sdk.VaultSdkSource
 import com.x8bit.bitwarden.ui.platform.feature.settings.appearance.model.AppLanguage
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -12,6 +13,7 @@ import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.stateIn
+import kotlinx.coroutines.launch
 
 /**
  * Primary implementation of [SettingsRepository].
@@ -20,6 +22,7 @@ import kotlinx.coroutines.flow.stateIn
 class SettingsRepositoryImpl(
     private val authDiskSource: AuthDiskSource,
     private val settingsDiskSource: SettingsDiskSource,
+    private val vaultSdkSource: VaultSdkSource,
     private val dispatcherManager: DispatcherManager,
 ) : SettingsRepository {
     private val activeUserId: String? get() = authDiskSource.userState?.activeUserId
@@ -77,6 +80,10 @@ class SettingsRepositoryImpl(
                 vaultTimeoutAction = value,
             )
         }
+    override val isUnlockWithPinEnabled: Boolean
+        get() = activeUserId
+            ?.let { authDiskSource.getEncryptedPin(userId = it) != null }
+            ?: false
 
     override fun clearData(userId: String) {
         settingsDiskSource.clearData(userId = userId)
@@ -159,6 +166,54 @@ class SettingsRepositoryImpl(
             )
         }
     }
+
+    override fun storeUnlockPin(
+        pin: String,
+        shouldRequireMasterPasswordOnRestart: Boolean,
+    ) {
+        val userId = activeUserId ?: return
+        unconfinedScope.launch {
+            vaultSdkSource
+                .derivePinKey(
+                    userId = userId,
+                    pin = pin,
+                )
+                .fold(
+                    onSuccess = { derivePinKeyResponse ->
+                        authDiskSource.apply {
+                            storeEncryptedPin(
+                                userId = userId,
+                                encryptedPin = derivePinKeyResponse.encryptedPin,
+                            )
+                            storePinProtectedUserKey(
+                                userId = userId,
+                                pinProtectedUserKey = derivePinKeyResponse.pinProtectedUserKey,
+                                inMemoryOnly = shouldRequireMasterPasswordOnRestart,
+                            )
+                        }
+                    },
+                    onFailure = {
+                        // PIN derivation should only fail when the user's vault is locked. This
+                        // should not be a concern when this method is actually called so we should
+                        // be able to safely ignore this.
+                    },
+                )
+        }
+    }
+
+    override fun clearUnlockPin() {
+        val userId = activeUserId ?: return
+        authDiskSource.apply {
+            storeEncryptedPin(
+                userId = userId,
+                encryptedPin = null,
+            )
+            authDiskSource.storePinProtectedUserKey(
+                userId = userId,
+                pinProtectedUserKey = null,
+            )
+        }
+    }
 }
 
 /**

app/src/main/java/com/x8bit/bitwarden/data/platform/repository/di/PlatformRepositoryModule.kt

@@ -8,6 +8,7 @@ import com.x8bit.bitwarden.data.platform.repository.EnvironmentRepository
 import com.x8bit.bitwarden.data.platform.repository.EnvironmentRepositoryImpl
 import com.x8bit.bitwarden.data.platform.repository.SettingsRepository
 import com.x8bit.bitwarden.data.platform.repository.SettingsRepositoryImpl
+import com.x8bit.bitwarden.data.vault.datasource.sdk.VaultSdkSource
 import dagger.Module
 import dagger.Provides
 import dagger.hilt.InstallIn
@@ -39,11 +40,13 @@ object PlatformRepositoryModule {
     fun provideSettingsRepository(
         authDiskSource: AuthDiskSource,
         settingsDiskSource: SettingsDiskSource,
+        vaultSdkSource: VaultSdkSource,
         dispatcherManager: DispatcherManager,
     ): SettingsRepository =
         SettingsRepositoryImpl(
             authDiskSource = authDiskSource,
             settingsDiskSource = settingsDiskSource,
+            vaultSdkSource = vaultSdkSource,
             dispatcherManager = dispatcherManager,
         )
 }

app/src/main/java/com/x8bit/bitwarden/ui/platform/feature/settings/accountsecurity/AccountSecurityViewModel.kt

@@ -38,7 +38,7 @@ class AccountSecurityViewModel @Inject constructor(
             fingerprintPhrase = "fingerprint-placeholder".asText(),
             isApproveLoginRequestsEnabled = false,
             isUnlockWithBiometricsEnabled = false,
-            isUnlockWithPinEnabled = false,
+            isUnlockWithPinEnabled = settingsRepository.isUnlockWithPinEnabled,
             vaultTimeout = settingsRepository.vaultTimeout,
             vaultTimeoutAction = settingsRepository.vaultTimeoutAction,
         ),
@@ -194,9 +194,17 @@ class AccountSecurityViewModel @Inject constructor(
         // TODO: Complete implementation (BIT-465)
         when (action) {
             AccountSecurityAction.UnlockWithPinToggle.PendingEnabled -> Unit
-            AccountSecurityAction.UnlockWithPinToggle.Disabled,
+            AccountSecurityAction.UnlockWithPinToggle.Disabled -> {
-            is AccountSecurityAction.UnlockWithPinToggle.Enabled,
+                settingsRepository.clearUnlockPin()
-            -> {
+                sendEvent(AccountSecurityEvent.ShowToast("Handle unlock with pin.".asText()))
+            }
+
+            is AccountSecurityAction.UnlockWithPinToggle.Enabled -> {
+                settingsRepository.storeUnlockPin(
+                    pin = action.pin,
+                    shouldRequireMasterPasswordOnRestart =
+                    action.shouldRequireMasterPasswordOnRestart,
+                )
                 sendEvent(AccountSecurityEvent.ShowToast("Handle unlock with pin.".asText()))
             }
         }

app/src/test/java/com/x8bit/bitwarden/data/auth/datasource/disk/util/FakeAuthDiskSource.kt

@@ -23,7 +23,7 @@ class FakeAuthDiskSource : AuthDiskSource {
     private val storedUserKeys = mutableMapOf<String, String?>()
     private val storedPrivateKeys = mutableMapOf<String, String?>()
     private val storedUserAutoUnlockKeys = mutableMapOf<String, String?>()
-    private val storedPinProtectedUserKeys = mutableMapOf<String, String?>()
+    private val storedPinProtectedUserKeys = mutableMapOf<String, Pair<String?, Boolean>>()
     private val storedEncryptedPins = mutableMapOf<String, String?>()
     private val storedOrganizations =
         mutableMapOf<String, List<SyncResponseJson.Profile.Organization>?>()
@@ -81,10 +81,14 @@ class FakeAuthDiskSource : AuthDiskSource {
     }
 
     override fun getPinProtectedUserKey(userId: String): String? =
-        storedPinProtectedUserKeys[userId]
+        storedPinProtectedUserKeys[userId]?.first
 
-    override fun storePinProtectedUserKey(userId: String, pinProtectedUserKey: String?) {
+    override fun storePinProtectedUserKey(
-        storedPinProtectedUserKeys[userId] = pinProtectedUserKey
+        userId: String,
+        pinProtectedUserKey: String?,
+        isInMemoryOnly: Boolean,
+    ) {
+        storedPinProtectedUserKeys[userId] = pinProtectedUserKey to isInMemoryOnly
     }
 
     override fun getEncryptedPin(userId: String): String? =
@@ -157,6 +161,24 @@ class FakeAuthDiskSource : AuthDiskSource {
         assertEquals(userAutoUnlockKey, storedUserAutoUnlockKeys[userId])
     }
 
+    /**
+     * Assert that the [encryptedPin] was stored successfully using the [userId].
+     */
+    fun assertEncryptedPin(userId: String, encryptedPin: String?) {
+        assertEquals(encryptedPin, storedEncryptedPins[userId])
+    }
+
+    /**
+     * Assert that the [pinProtectedUserKey] was stored successfully using the [userId].
+     */
+    fun assertPinProtectedUserKey(
+        userId: String,
+        pinProtectedUserKey: String?,
+        inMemoryOnly: Boolean = false,
+    ) {
+        assertEquals(pinProtectedUserKey to inMemoryOnly, storedPinProtectedUserKeys[userId])
+    }
+
     /**
      * Assert the the [organizationKeys] was stored successfully using the [userId].
      */

app/src/test/java/com/x8bit/bitwarden/data/platform/repository/SettingsRepositoryTest.kt

@@ -1,15 +1,18 @@
 package com.x8bit.bitwarden.data.platform.repository
 
 import app.cash.turbine.test
-import com.x8bit.bitwarden.data.auth.datasource.disk.AuthDiskSource
+import com.bitwarden.core.DerivePinKeyResponse
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.UserStateJson
+import com.x8bit.bitwarden.data.auth.datasource.disk.util.FakeAuthDiskSource
 import com.x8bit.bitwarden.data.platform.base.FakeDispatcherManager
 import com.x8bit.bitwarden.data.platform.datasource.disk.util.FakeSettingsDiskSource
 import com.x8bit.bitwarden.data.platform.repository.model.VaultTimeout
 import com.x8bit.bitwarden.data.platform.repository.model.VaultTimeoutAction
+import com.x8bit.bitwarden.data.platform.util.asSuccess
+import com.x8bit.bitwarden.data.vault.datasource.sdk.VaultSdkSource
 import com.x8bit.bitwarden.ui.platform.feature.settings.appearance.model.AppLanguage
 import io.mockk.coEvery
-import io.mockk.every
+import io.mockk.coVerify
 import io.mockk.mockk
 import kotlinx.coroutines.test.runTest
 import org.junit.jupiter.api.Assertions.assertEquals
@@ -19,12 +22,14 @@ import org.junit.jupiter.api.Assertions.assertTrue
 import org.junit.jupiter.api.Test
 
 class SettingsRepositoryTest {
-    private val authDiskSource: AuthDiskSource = mockk()
+    private val fakeAuthDiskSource = FakeAuthDiskSource()
     private val fakeSettingsDiskSource = FakeSettingsDiskSource()
+    private val vaultSdkSource: VaultSdkSource = mockk()
 
     private val settingsRepository = SettingsRepositoryImpl(
-        authDiskSource = authDiskSource,
+        authDiskSource = fakeAuthDiskSource,
         settingsDiskSource = fakeSettingsDiskSource,
+        vaultSdkSource = vaultSdkSource,
         dispatcherManager = FakeDispatcherManager(),
     )
 
@@ -127,14 +132,14 @@ class SettingsRepositoryTest {
 
     @Test
     fun `vaultTimeout should pull from and update SettingsDiskSource for the current user`() {
-        every { authDiskSource.userState?.activeUserId } returns null
+        fakeAuthDiskSource.userState = null
         assertEquals(
             VaultTimeout.Never,
             settingsRepository.vaultTimeout,
         )
 
         val userId = "userId"
-        every { authDiskSource.userState?.activeUserId } returns userId
+        fakeAuthDiskSource.userState = MOCK_USER_STATE
 
         // Updates to the disk source change the repository value
         VAULT_TIMEOUT_MAP.forEach { (vaultTimeout, vaultTimeoutInMinutes) ->
@@ -160,14 +165,14 @@ class SettingsRepositoryTest {
 
     @Test
     fun `vaultTimeoutAction should pull from and update SettingsDiskSource`() {
-        every { authDiskSource.userState?.activeUserId } returns null
+        fakeAuthDiskSource.userState = null
         assertEquals(
             VaultTimeoutAction.LOCK,
             settingsRepository.vaultTimeoutAction,
         )
 
         val userId = "userId"
-        every { authDiskSource.userState?.activeUserId } returns userId
+        fakeAuthDiskSource.userState = MOCK_USER_STATE
 
         // Updates to the disk source change the repository value
         VAULT_TIMEOUT_ACTIONS.forEach { vaultTimeoutAction ->
@@ -287,13 +292,28 @@ class SettingsRepositoryTest {
         }
     }
 
+    @Suppress("MaxLineLength")
+    @Test
+    fun `isUnlockWithPinEnabled should return a value that tracks the existence of an encrypted PIN for the current user`() {
+        val userId = "userId"
+        fakeAuthDiskSource.userState = MOCK_USER_STATE
+        fakeAuthDiskSource.storeEncryptedPin(
+            userId = userId,
+            encryptedPin = null,
+        )
+        assertFalse(settingsRepository.isUnlockWithPinEnabled)
+
+        fakeAuthDiskSource.storeEncryptedPin(
+            userId = userId,
+            encryptedPin = "encryptedPin",
+        )
+        assertTrue(settingsRepository.isUnlockWithPinEnabled)
+    }
+
     @Test
     fun `getPullToRefreshEnabledFlow should react to changes in SettingsDiskSource`() = runTest {
         val userId = "userId"
-        val userState = mockk<UserStateJson> {
+        fakeAuthDiskSource.userState = MOCK_USER_STATE
-            every { activeUserId } returns userId
-        }
-        coEvery { authDiskSource.userState } returns userState
         settingsRepository
             .getPullToRefreshEnabledFlow()
             .test {
@@ -314,12 +334,134 @@ class SettingsRepositoryTest {
     @Test
     fun `storePullToRefreshEnabled should properly update SettingsDiskSource`() {
         val userId = "userId"
-        every { authDiskSource.userState?.activeUserId } returns userId
+        fakeAuthDiskSource.userState = MOCK_USER_STATE
         settingsRepository.storePullToRefreshEnabled(true)
         assertEquals(true, fakeSettingsDiskSource.getPullToRefreshEnabled(userId = userId))
     }
+
+    @Suppress("MaxLineLength")
+    @Test
+    fun `storeUnlockPin when the master password on restart is required should only save an encrypted PIN to disk`() {
+        val userId = "userId"
+        val pin = "1234"
+        val encryptedPin = "encryptedPin"
+        val pinProtectedUserKey = "pinProtectedUserKey"
+        val derivePinKeyResponse = DerivePinKeyResponse(
+            pinProtectedUserKey = pinProtectedUserKey,
+            encryptedPin = encryptedPin,
+        )
+        fakeAuthDiskSource.userState = MOCK_USER_STATE
+        coEvery {
+            vaultSdkSource.derivePinKey(
+                userId = userId,
+                pin = pin,
+            )
+        } returns derivePinKeyResponse.asSuccess()
+
+        settingsRepository.storeUnlockPin(
+            pin = pin,
+            shouldRequireMasterPasswordOnRestart = true,
+        )
+
+        fakeAuthDiskSource.apply {
+            assertEncryptedPin(
+                userId = userId,
+                encryptedPin = encryptedPin,
+            )
+            assertPinProtectedUserKey(
+                userId = userId,
+                pinProtectedUserKey = pinProtectedUserKey,
+                inMemoryOnly = true,
+            )
+        }
+        coVerify {
+            vaultSdkSource.derivePinKey(
+                userId = userId,
+                pin = pin,
+            )
+        }
+    }
+
+    @Suppress("MaxLineLength")
+    @Test
+    fun `storeUnlockPin when the master password on restart is not required should save all PIN data to disk`() {
+        val userId = "userId"
+        val pin = "1234"
+        val encryptedPin = "encryptedPin"
+        val pinProtectedUserKey = "pinProtectedUserKey"
+        val derivePinKeyResponse = DerivePinKeyResponse(
+            pinProtectedUserKey = pinProtectedUserKey,
+            encryptedPin = encryptedPin,
+        )
+        fakeAuthDiskSource.userState = MOCK_USER_STATE
+        coEvery {
+            vaultSdkSource.derivePinKey(
+                userId = userId,
+                pin = pin,
+            )
+        } returns derivePinKeyResponse.asSuccess()
+
+        settingsRepository.storeUnlockPin(
+            pin = pin,
+            shouldRequireMasterPasswordOnRestart = false,
+        )
+
+        fakeAuthDiskSource.apply {
+            assertEncryptedPin(
+                userId = userId,
+                encryptedPin = encryptedPin,
+            )
+            assertPinProtectedUserKey(
+                userId = userId,
+                pinProtectedUserKey = pinProtectedUserKey,
+                inMemoryOnly = false,
+            )
+        }
+        coVerify {
+            vaultSdkSource.derivePinKey(
+                userId = userId,
+                pin = pin,
+            )
+        }
+    }
+
+    @Suppress("MaxLineLength")
+    @Test
+    fun `clearUnlockPin should clear any previously stored PIN-related values for the current user`() {
+        val userId = "userId"
+        fakeAuthDiskSource.userState = MOCK_USER_STATE
+        fakeAuthDiskSource.apply {
+            storeEncryptedPin(
+                userId = userId,
+                encryptedPin = "encryptedPin",
+            )
+            storePinProtectedUserKey(
+                userId = userId,
+                pinProtectedUserKey = "pinProtectedUserKey",
+            )
+        }
+
+        settingsRepository.clearUnlockPin()
+
+        fakeAuthDiskSource.apply {
+            assertEncryptedPin(
+                userId = userId,
+                encryptedPin = null,
+            )
+            assertPinProtectedUserKey(
+                userId = userId,
+                pinProtectedUserKey = null,
+            )
+        }
+    }
 }
 
+private val MOCK_USER_STATE =
+    UserStateJson(
+        activeUserId = "userId",
+        accounts = mapOf("userId" to mockk()),
+    )
+
 /**
  * A list of all [VaultTimeoutAction].
  *

app/src/test/java/com/x8bit/bitwarden/ui/platform/feature/settings/accountsecurity/AccountSecurityViewModelTest.kt

@@ -22,11 +22,28 @@ import org.junit.jupiter.api.Test
 class AccountSecurityViewModelTest : BaseViewModelTest() {
 
     @Test
-    fun `initial state should be correct`() {
+    fun `initial state should be correct when saved state is set`() {
-        val viewModel = createViewModel()
+        val viewModel = createViewModel(initialState = DEFAULT_STATE)
         assertEquals(DEFAULT_STATE, viewModel.stateFlow.value)
     }
 
+    @Test
+    fun `initial state should be correct when saved state is not set`() {
+        val settingsRepository: SettingsRepository = mockk {
+            every { isUnlockWithPinEnabled } returns true
+            every { vaultTimeout } returns VaultTimeout.ThirtyMinutes
+            every { vaultTimeoutAction } returns VaultTimeoutAction.LOCK
+        }
+        val viewModel = createViewModel(
+            initialState = null,
+            settingsRepository = settingsRepository,
+        )
+        assertEquals(
+            DEFAULT_STATE.copy(isUnlockWithPinEnabled = true),
+            viewModel.stateFlow.value,
+        )
+    }
+
     @Test
     fun `on AccountFingerprintPhraseClick should show the fingerprint phrase dialog`() = runTest {
         val viewModel = createViewModel()
@@ -210,12 +227,18 @@ class AccountSecurityViewModelTest : BaseViewModelTest() {
 
     @Suppress("MaxLineLength")
     @Test
-    fun `on UnlockWithPinToggle Disabled should set pin unlock to false and emit ShowToast`() =
+    fun `on UnlockWithPinToggle Disabled should set pin unlock to false, clear the PIN in settings, and emit ShowToast`() =
         runTest {
             val initialState = DEFAULT_STATE.copy(
                 isUnlockWithPinEnabled = true,
             )
-            val viewModel = createViewModel(initialState = initialState)
+            val settingsRepository: SettingsRepository = mockk() {
+                every { clearUnlockPin() } just runs
+            }
+            val viewModel = createViewModel(
+                initialState = initialState,
+                settingsRepository = settingsRepository,
+            )
             viewModel.eventFlow.test {
                 viewModel.trySendAction(
                     AccountSecurityAction.UnlockWithPinToggle.Disabled,
@@ -229,11 +252,12 @@ class AccountSecurityViewModelTest : BaseViewModelTest() {
                 initialState.copy(isUnlockWithPinEnabled = false),
                 viewModel.stateFlow.value,
             )
+            verify { settingsRepository.clearUnlockPin() }
         }
 
     @Suppress("MaxLineLength")
     @Test
-    fun `on UnlockWithPinToggle Enabled should set pin unlock to true`() {
+    fun `on UnlockWithPinToggle PendingEnabled should set pin unlock to true`() {
         val initialState = DEFAULT_STATE.copy(
             isUnlockWithPinEnabled = false,
         )
@@ -249,12 +273,18 @@ class AccountSecurityViewModelTest : BaseViewModelTest() {
 
     @Suppress("MaxLineLength")
     @Test
-    fun `on UnlockWithPinToggle Enabled should set pin unlock to true and emit ShowToast`() =
+    fun `on UnlockWithPinToggle Enabled should set pin unlock to true, set the PIN in settings, and emit ShowToast`() =
         runTest {
             val initialState = DEFAULT_STATE.copy(
                 isUnlockWithPinEnabled = false,
             )
-            val viewModel = createViewModel(initialState = initialState)
+            val settingsRepository: SettingsRepository = mockk() {
+                every { storeUnlockPin(any(), any()) } just runs
+            }
+            val viewModel = createViewModel(
+                initialState = initialState,
+                settingsRepository = settingsRepository,
+            )
             viewModel.eventFlow.test {
                 viewModel.trySendAction(
                     AccountSecurityAction.UnlockWithPinToggle.Enabled(
@@ -271,6 +301,12 @@ class AccountSecurityViewModelTest : BaseViewModelTest() {
                 initialState.copy(isUnlockWithPinEnabled = true),
                 viewModel.stateFlow.value,
             )
+            verify {
+                settingsRepository.storeUnlockPin(
+                    pin = "1234",
+                    shouldRequireMasterPasswordOnRestart = true,
+                )
+            }
         }
 
     @Test
@@ -302,7 +338,7 @@ class AccountSecurityViewModelTest : BaseViewModelTest() {
     }
 
     private fun createViewModel(
-        initialState: AccountSecurityState = DEFAULT_STATE,
+        initialState: AccountSecurityState? = DEFAULT_STATE,
         authRepository: AuthRepository = mockk(relaxed = true),
         vaultRepository: VaultRepository = mockk(relaxed = true),
         settingsRepository: SettingsRepository = mockk(relaxed = true),