refactored personal vault migration logic into VaultMigrationManager

2026-10-23 03:23:53 -05:00 · 2026-01-06 12:03:43 +00:00
1 parent 3247cd12b1
commit 23737ef18e
@@ -0,0 +1,34 @@
+package com.x8bit.bitwarden.data.vault.manager
+
+import com.bitwarden.network.model.SyncResponseJson
+import com.x8bit.bitwarden.data.vault.manager.model.VaultMigrationData
+import kotlinx.coroutines.flow.StateFlow
+
+/**
+ * Manages the migration of personal vault items to organization collections.
+ * This interface provides a way to check if migration is needed and track migration state.
+ */
+interface VaultMigrationManager {
+    /**
+     * Flow that emits when conditions are met for the user to migrate their personal vault.
+     * Updated after each sync to reflect current policy and vault state.
+     */
+    val vaultMigrationDataStateFlow: StateFlow<VaultMigrationData>
+
+    /**
+     * Verifies if the user should migrate their personal vault to organization collections
+     * based on active policies, feature flags, and the provided cipher list.
+     *
+     * @param cipherList List of ciphers from the sync response to check for personal items.
+     */
+    fun verifyAndUpdateMigrationState(cipherList: List<SyncResponseJson.Cipher>)
+
+    /**
+     * Checks if the user should migrate their vault based on policies, feature flags,
+     * network connectivity, and whether they have personal items.
+     *
+     * @param hasPersonalItems Callback to check if the user has personal items.
+     * @return true if migration conditions are met, false otherwise.
+     */
+    fun shouldMigrateVault(hasPersonalItems: () -> Boolean): Boolean
+}
@@ -0,0 +1,67 @@
+package com.x8bit.bitwarden.data.vault.manager
+
+import com.bitwarden.core.data.manager.model.FlagKey
+import com.bitwarden.network.model.PolicyTypeJson
+import com.bitwarden.network.model.SyncResponseJson
+import com.x8bit.bitwarden.data.auth.datasource.disk.AuthDiskSource
+import com.x8bit.bitwarden.data.platform.manager.FeatureFlagManager
+import com.x8bit.bitwarden.data.platform.manager.PolicyManager
+import com.x8bit.bitwarden.data.platform.manager.network.NetworkConnectionManager
+import com.x8bit.bitwarden.data.vault.manager.model.VaultMigrationData
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.flow.update
+
+/**
+ * Default implementation of [VaultMigrationManager].
+ */
+class VaultMigrationManagerImpl(
+    private val authDiskSource: AuthDiskSource,
+    private val policyManager: PolicyManager,
+    private val featureFlagManager: FeatureFlagManager,
+    private val connectionManager: NetworkConnectionManager,
+) : VaultMigrationManager {
+
+    private val activeUserId: String? get() = authDiskSource.userState?.activeUserId
+
+    private val mutableVaultMigrationDataStateFlow =
+        MutableStateFlow<VaultMigrationData>(value = VaultMigrationData.NoMigrationRequired)
+
+    override val vaultMigrationDataStateFlow: StateFlow<VaultMigrationData>
+        get() = mutableVaultMigrationDataStateFlow.asStateFlow()
+
+    override fun verifyAndUpdateMigrationState(cipherList: List<SyncResponseJson.Cipher>) {
+        val userId = activeUserId ?: return
+
+        mutableVaultMigrationDataStateFlow.update {
+            if (shouldMigrateVault { cipherList.any { it.organizationId == null } }) {
+                val orgId = policyManager.getPersonalOwnershipPolicyOrganizationId()
+                val orgName = authDiskSource
+                    .getOrganizations(userId = userId)
+                    ?.firstOrNull { it.id == orgId }
+                    ?.name
+
+                if (orgId != null && orgName != null) {
+                    VaultMigrationData.MigrationRequired(
+                        organizationId = orgId,
+                        organizationName = orgName,
+                    )
+                } else {
+                    VaultMigrationData.NoMigrationRequired
+                }
+            } else {
+                VaultMigrationData.NoMigrationRequired
+            }
+        }
+    }
+
+    override fun shouldMigrateVault(hasPersonalItems: () -> Boolean): Boolean {
+        return policyManager
+            .getActivePolicies(PolicyTypeJson.PERSONAL_OWNERSHIP)
+            .any() &&
+            featureFlagManager.getFeatureFlag(FlagKey.MigrateMyVaultToMyItems) &&
+            connectionManager.isNetworkConnected &&
+            hasPersonalItems()
+    }
+}
@@ -5,7 +5,6 @@ import com.bitwarden.core.data.repository.model.DataState
 import com.bitwarden.vault.DecryptCipherListResult
 import com.bitwarden.vault.FolderView
 import com.x8bit.bitwarden.data.vault.manager.model.SyncVaultDataResult
-import com.x8bit.bitwarden.data.vault.manager.model.VaultMigrationData
 import com.x8bit.bitwarden.data.vault.repository.model.DomainsData
 import com.x8bit.bitwarden.data.vault.repository.model.SendData
 import com.x8bit.bitwarden.data.vault.repository.model.VaultData
@@ -63,12 +62,6 @@ interface VaultSyncManager {
     */
    val sendDataStateFlow: StateFlow<DataState<SendData>>

-    /**
-     * Flow that emits when conditions are met for the user to migrate their personal vault.
-     * Updated after each sync to reflect current policy and vault state.
-     */
-    val shouldMigratePersonalVaultFlow: StateFlow<VaultMigrationData>
-
    /**
     * Sync the vault data for the current user.
     *
@@ -3,14 +3,11 @@ package com.x8bit.bitwarden.data.vault.manager
 import com.bitwarden.collections.CollectionView
 import com.bitwarden.core.InitOrgCryptoRequest
 import com.bitwarden.core.data.manager.dispatcher.DispatcherManager
-import com.bitwarden.core.data.manager.model.FlagKey
 import com.bitwarden.core.data.repository.model.DataState
 import com.bitwarden.core.data.repository.util.combineDataStates
 import com.bitwarden.core.data.repository.util.map
 import com.bitwarden.core.data.repository.util.updateToPendingOrLoading
-import com.bitwarden.network.model.PolicyTypeJson
 import com.bitwarden.network.model.SyncResponseJson
-import com.bitwarden.network.model.SyncResponseJson.Cipher
 import com.bitwarden.network.service.SyncService
 import com.bitwarden.network.util.isNoConnectionError
 import com.bitwarden.vault.DecryptCipherListResult
@@ -25,16 +22,12 @@ import com.x8bit.bitwarden.data.platform.datasource.disk.SettingsDiskSource
 import com.x8bit.bitwarden.data.platform.error.NoActiveUserException
 import com.x8bit.bitwarden.data.platform.error.SecurityStampMismatchException
 import com.x8bit.bitwarden.data.platform.manager.DatabaseSchemeManager
-import com.x8bit.bitwarden.data.platform.manager.FeatureFlagManager
-import com.x8bit.bitwarden.data.platform.manager.PolicyManager
 import com.x8bit.bitwarden.data.platform.manager.PushManager
-import com.x8bit.bitwarden.data.platform.manager.network.NetworkConnectionManager
 import com.x8bit.bitwarden.data.platform.repository.util.observeWhenSubscribedAndLoggedIn
 import com.x8bit.bitwarden.data.platform.repository.util.observeWhenSubscribedAndUnlocked
 import com.x8bit.bitwarden.data.vault.datasource.disk.VaultDiskSource
 import com.x8bit.bitwarden.data.vault.datasource.sdk.VaultSdkSource
 import com.x8bit.bitwarden.data.vault.manager.model.SyncVaultDataResult
-import com.x8bit.bitwarden.data.vault.manager.model.VaultMigrationData
 import com.x8bit.bitwarden.data.vault.repository.model.DomainsData
 import com.x8bit.bitwarden.data.vault.repository.model.SendData
 import com.x8bit.bitwarden.data.vault.repository.model.VaultData
@@ -88,9 +81,7 @@ class VaultSyncManagerImpl(
    private val userLogoutManager: UserLogoutManager,
    private val userStateManager: UserStateManager,
    private val vaultLockManager: VaultLockManager,
-    private val policyManager: PolicyManager,
-    private val featureFlagManager: FeatureFlagManager,
-    private val connectionManager: NetworkConnectionManager,
+    private val vaultMigrationManager: VaultMigrationManager,
    private val clock: Clock,
    databaseSchemeManager: DatabaseSchemeManager,
    pushManager: PushManager,
@@ -117,9 +108,6 @@ class VaultSyncManagerImpl(
    private val mutableDomainsStateFlow =
        MutableStateFlow<DataState<DomainsData>>(DataState.Loading)

-    private val mutableShouldMigratePersonalVaultFlow =
-        MutableStateFlow<VaultMigrationData>(value = VaultMigrationData.NoMigrationRequired)
-
    override val decryptCipherListResultStateFlow: StateFlow<DataState<DecryptCipherListResult>>
        get() = mutableDecryptCipherListResultFlow.asStateFlow()

@@ -135,9 +123,6 @@ class VaultSyncManagerImpl(
    override val sendDataStateFlow: StateFlow<DataState<SendData>>
        get() = mutableSendDataStateFlow.asStateFlow()

-    override val shouldMigratePersonalVaultFlow: StateFlow<VaultMigrationData>
-        get() = mutableShouldMigratePersonalVaultFlow.asStateFlow()
-
    override val vaultDataStateFlow: StateFlow<DataState<VaultData>> =
        combine(
            decryptCipherListResultStateFlow,
@@ -358,7 +343,7 @@ class VaultSyncManagerImpl(
                    vaultDiskSource.replaceVaultData(userId = userId, vault = syncResponse)
                    val itemsAvailable = syncResponse.ciphers?.isNotEmpty() == true
                    syncResponse.ciphers?.let {
-                        verifyAndUpdateIfUserShouldMigrateVaultToMyItems(it)
+                        vaultMigrationManager.verifyAndUpdateMigrationState(it)
                    }
                    SyncVaultDataResult.Success(itemsAvailable = itemsAvailable)
                }
@@ -430,7 +415,7 @@ class VaultSyncManagerImpl(
                        onSuccess = { result ->
                            // We need to be sure the data on device is updated
                            // before sending the user to the migration screen
-                            if (userShouldMigrateVault {
+                            if (vaultMigrationManager.shouldMigrateVault {
                                    result.successes.any { it.organizationId == null }
                                }
                            ) {
@@ -550,44 +535,6 @@ class VaultSyncManagerImpl(
        }
    }

-    /**
-     * Evaluates if the user should migrate their personal vault based on
-     * policies, personal ciphers, feature flag, and network connectivity.
-     */
-    private fun userShouldMigrateVault(hasPersonalItems: () -> Boolean): Boolean {
-        return policyManager
-            .getActivePolicies(PolicyTypeJson.PERSONAL_OWNERSHIP)
-            .any() &&
-            featureFlagManager.getFeatureFlag(FlagKey.MigrateMyVaultToMyItems) &&
-            connectionManager.isNetworkConnected &&
-            hasPersonalItems()
-    }
-
-    private fun verifyAndUpdateIfUserShouldMigrateVaultToMyItems(cipherList: List<Cipher>) {
-        val userId = activeUserId ?: return
-
-        mutableShouldMigratePersonalVaultFlow.update {
-            if (userShouldMigrateVault { cipherList.any { it.organizationId == null } }) {
-                val orgId = policyManager.getPersonalOwnershipPolicyOrganizationId()
-                val orgName = authDiskSource
-                    .getOrganizations(userId = userId)
-                    ?.firstOrNull { it.id == orgId }
-                    ?.name
-
-                if (orgId != null && orgName != null) {
-                    VaultMigrationData.MigrationRequired(
-                        organizationId = orgId,
-                        organizationName = orgName,
-                    )
-                } else {
-                    VaultMigrationData.NoMigrationRequired
-                }
-            } else {
-                VaultMigrationData.NoMigrationRequired
-            }
-        }
-    }
-
    /**
     * Returns the given [DataState] as-is, or [DataState.Loading] if vault data for the given
     * [userId] has not synced. This can be used to distinguish between empty data in the database
@@ -39,6 +39,8 @@ import com.x8bit.bitwarden.data.vault.manager.TotpCodeManager
 import com.x8bit.bitwarden.data.vault.manager.TotpCodeManagerImpl
 import com.x8bit.bitwarden.data.vault.manager.VaultLockManager
 import com.x8bit.bitwarden.data.vault.manager.VaultLockManagerImpl
+import com.x8bit.bitwarden.data.vault.manager.VaultMigrationManager
+import com.x8bit.bitwarden.data.vault.manager.VaultMigrationManagerImpl
 import com.x8bit.bitwarden.data.vault.manager.VaultSyncManager
 import com.x8bit.bitwarden.data.vault.manager.VaultSyncManagerImpl
 import dagger.Module
@@ -57,6 +59,20 @@ import javax.inject.Singleton
@InstallIn(SingletonComponent::class)
 object VaultManagerModule {

+    @Provides
+    @Singleton
+    fun provideVaultMigrationManager(
+        authDiskSource: AuthDiskSource,
+        policyManager: PolicyManager,
+        featureFlagManager: FeatureFlagManager,
+        connectionManager: NetworkConnectionManager,
+    ): VaultMigrationManager = VaultMigrationManagerImpl(
+        authDiskSource = authDiskSource,
+        policyManager = policyManager,
+        featureFlagManager = featureFlagManager,
+        connectionManager = connectionManager,
+    )
+
    @Provides
    @Singleton
    fun provideCipherManager(
@@ -195,9 +211,7 @@ object VaultManagerModule {
        userLogoutManager: UserLogoutManager,
        userStateManager: UserStateManager,
        vaultLockManager: VaultLockManager,
-        policyManager: PolicyManager,
-        featureFlagManager: FeatureFlagManager,
-        connectionManager: NetworkConnectionManager,
+        vaultMigrationManager: VaultMigrationManager,
        clock: Clock,
        databaseSchemeManager: DatabaseSchemeManager,
        pushManager: PushManager,
@@ -211,9 +225,7 @@ object VaultManagerModule {
        userLogoutManager = userLogoutManager,
        userStateManager = userStateManager,
        vaultLockManager = vaultLockManager,
-        policyManager = policyManager,
-        featureFlagManager = featureFlagManager,
-        connectionManager = connectionManager,
+        vaultMigrationManager = vaultMigrationManager,
        clock = clock,
        databaseSchemeManager = databaseSchemeManager,
        pushManager = pushManager,
@@ -19,8 +19,8 @@ import com.x8bit.bitwarden.data.credentials.model.GetCredentialsRequest
 import com.x8bit.bitwarden.data.credentials.model.ProviderGetPasswordCredentialRequest
 import com.x8bit.bitwarden.data.platform.manager.SpecialCircumstanceManager
 import com.x8bit.bitwarden.data.platform.manager.model.SpecialCircumstance
+import com.x8bit.bitwarden.data.vault.manager.VaultMigrationManager
 import com.x8bit.bitwarden.data.vault.manager.model.VaultMigrationData
-import com.x8bit.bitwarden.data.vault.repository.VaultRepository
 import com.x8bit.bitwarden.ui.tools.feature.send.model.SendItemType
 import dagger.hilt.android.lifecycle.HiltViewModel
 import kotlinx.coroutines.flow.combine
@@ -37,7 +37,7 @@ import javax.inject.Inject
 class RootNavViewModel @Inject constructor(
    private val authRepository: AuthRepository,
    specialCircumstanceManager: SpecialCircumstanceManager,
-    vaultRepository: VaultRepository,
+    vaultMigrationManager: VaultMigrationManager,
 ) : BaseViewModel<RootNavState, Unit, RootNavAction>(
    initialState = RootNavState.Splash,
 ) {
@@ -46,7 +46,7 @@ class RootNavViewModel @Inject constructor(
            authRepository.authStateFlow,
            authRepository.userStateFlow,
            specialCircumstanceManager.specialCircumstanceStateFlow,
-            vaultRepository.shouldMigratePersonalVaultFlow,
+            vaultMigrationManager.vaultMigrationDataStateFlow,
        ) { authState, userState, specialCircumstance, shouldMigratePersonalVault ->
            RootNavAction.Internal.UserStateUpdateReceive(
                authState = authState,
@@ -0,0 +1,424 @@
+package com.x8bit.bitwarden.data.vault.manager
+
+import app.cash.turbine.test
+import com.bitwarden.core.data.manager.model.FlagKey
+import com.bitwarden.network.model.PolicyTypeJson
+import com.bitwarden.network.model.SyncResponseJson
+import com.bitwarden.network.model.createMockCipher
+import com.bitwarden.network.model.createMockPolicy
+import com.bitwarden.network.model.createMockProfile
+import com.bitwarden.network.model.createMockSyncResponse
+import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountJson
+import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountTokensJson
+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.manager.FeatureFlagManager
+import com.x8bit.bitwarden.data.platform.manager.PolicyManager
+import com.x8bit.bitwarden.data.platform.manager.network.NetworkConnectionManager
+import com.x8bit.bitwarden.data.vault.manager.model.VaultMigrationData
+import io.mockk.every
+import io.mockk.mockk
+import kotlinx.coroutines.test.runTest
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Test
+import java.time.ZonedDateTime
+
+class VaultMigrationManagerTest {
+
+    private val fakeAuthDiskSource = FakeAuthDiskSource()
+    private val policyManager: PolicyManager = mockk {
+        every { getActivePolicies(any()) } returns emptyList()
+    }
+    private val mockFeatureFlagManager = mockk<FeatureFlagManager>(relaxed = true) {
+        every { getFeatureFlag(FlagKey.MigrateMyVaultToMyItems) } returns true
+    }
+    private val connectionManager: NetworkConnectionManager = mockk {
+        every { isNetworkConnected } returns true
+    }
+
+    private val vaultMigrationManager: VaultMigrationManager = VaultMigrationManagerImpl(
+        authDiskSource = fakeAuthDiskSource,
+        policyManager = policyManager,
+        featureFlagManager = mockFeatureFlagManager,
+        connectionManager = connectionManager,
+    )
+
+    @Test
+    fun `vaultMigrationDataStateFlow should initially emit NoMigrationRequired`() = runTest {
+        assertEquals(
+            VaultMigrationData.NoMigrationRequired,
+            vaultMigrationManager.vaultMigrationDataStateFlow.value,
+        )
+    }
+
+    @Test
+    @Suppress("MaxLineLength")
+    fun `verifyAndUpdateMigrationState should emit MigrationRequired when all conditions are met`() =
+        runTest {
+            val userId = "mockId-1"
+            fakeAuthDiskSource.userState = MOCK_USER_STATE
+
+            // Setup conditions for migration
+            val mockPolicy = createMockPolicy(number = 1, type = PolicyTypeJson.PERSONAL_OWNERSHIP)
+            every {
+                policyManager.getActivePolicies(PolicyTypeJson.PERSONAL_OWNERSHIP)
+            } returns listOf(mockPolicy)
+            every {
+                policyManager.getPersonalOwnershipPolicyOrganizationId()
+            } returns "mockId-1"
+
+            // Store organizations
+            val syncResponse = createMockSyncResponse(number = 1)
+            fakeAuthDiskSource.storeOrganizations(
+                userId = userId,
+                organizations = syncResponse.profile.organizations,
+            )
+
+            // Setup sync with personal ciphers (organizationId = null)
+            val personalCipher = createMockCipher(number = 1).copy(organizationId = null)
+            val cipherList = listOf(personalCipher)
+
+            vaultMigrationManager.vaultMigrationDataStateFlow.test {
+                assertEquals(VaultMigrationData.NoMigrationRequired, awaitItem())
+
+                vaultMigrationManager.verifyAndUpdateMigrationState(cipherList)
+
+                assertEquals(
+                    VaultMigrationData.MigrationRequired(
+                        organizationId = "mockId-1",
+                        organizationName = "mockName-1",
+                    ),
+                    awaitItem(),
+                )
+            }
+        }
+
+    @Test
+    @Suppress("MaxLineLength")
+    fun `verifyAndUpdateMigrationState should emit NoMigrationRequired when no personal ownership policy`() =
+        runTest {
+            fakeAuthDiskSource.userState = MOCK_USER_STATE
+
+            every {
+                policyManager.getActivePolicies(PolicyTypeJson.PERSONAL_OWNERSHIP)
+            } returns emptyList()
+
+            val personalCipher = createMockCipher(number = 1).copy(organizationId = null)
+            val cipherList = listOf(personalCipher)
+
+            vaultMigrationManager.vaultMigrationDataStateFlow.test {
+                assertEquals(VaultMigrationData.NoMigrationRequired, awaitItem())
+
+                vaultMigrationManager.verifyAndUpdateMigrationState(cipherList)
+
+                // Should still be NoMigrationRequired since no policy
+                expectNoEvents()
+            }
+        }
+
+    @Test
+    @Suppress("MaxLineLength")
+    fun `verifyAndUpdateMigrationState should emit NoMigrationRequired when feature flag is disabled`() =
+        runTest {
+            fakeAuthDiskSource.userState = MOCK_USER_STATE
+
+            val mockPolicy = createMockPolicy(number = 1, type = PolicyTypeJson.PERSONAL_OWNERSHIP)
+            every {
+                policyManager.getActivePolicies(PolicyTypeJson.PERSONAL_OWNERSHIP)
+            } returns listOf(mockPolicy)
+
+            every {
+                mockFeatureFlagManager.getFeatureFlag(FlagKey.MigrateMyVaultToMyItems)
+            } returns false
+
+            val personalCipher = createMockCipher(number = 1).copy(organizationId = null)
+            val cipherList = listOf(personalCipher)
+
+            vaultMigrationManager.vaultMigrationDataStateFlow.test {
+                assertEquals(VaultMigrationData.NoMigrationRequired, awaitItem())
+
+                vaultMigrationManager.verifyAndUpdateMigrationState(cipherList)
+
+                // Should still be NoMigrationRequired since feature flag is disabled
+                expectNoEvents()
+            }
+        }
+
+    @Test
+    @Suppress("MaxLineLength")
+    fun `verifyAndUpdateMigrationState should emit NoMigrationRequired when no network connection`() =
+        runTest {
+            fakeAuthDiskSource.userState = MOCK_USER_STATE
+
+            val mockPolicy = createMockPolicy(number = 1, type = PolicyTypeJson.PERSONAL_OWNERSHIP)
+            every {
+                policyManager.getActivePolicies(PolicyTypeJson.PERSONAL_OWNERSHIP)
+            } returns listOf(mockPolicy)
+
+            every { connectionManager.isNetworkConnected } returns false
+
+            val personalCipher = createMockCipher(number = 1).copy(organizationId = null)
+            val cipherList = listOf(personalCipher)
+
+            vaultMigrationManager.vaultMigrationDataStateFlow.test {
+                assertEquals(VaultMigrationData.NoMigrationRequired, awaitItem())
+
+                vaultMigrationManager.verifyAndUpdateMigrationState(cipherList)
+
+                // Should still be NoMigrationRequired since no network
+                expectNoEvents()
+            }
+        }
+
+    @Test
+    @Suppress("MaxLineLength")
+    fun `verifyAndUpdateMigrationState should emit NoMigrationRequired when no personal ciphers exist`() =
+        runTest {
+            fakeAuthDiskSource.userState = MOCK_USER_STATE
+
+            val mockPolicy = createMockPolicy(number = 1, type = PolicyTypeJson.PERSONAL_OWNERSHIP)
+            every {
+                policyManager.getActivePolicies(PolicyTypeJson.PERSONAL_OWNERSHIP)
+            } returns listOf(mockPolicy)
+
+            // All ciphers belong to organizations
+            val orgCipher = createMockCipher(number = 1).copy(organizationId = "org-id")
+            val cipherList = listOf(orgCipher)
+
+            vaultMigrationManager.vaultMigrationDataStateFlow.test {
+                assertEquals(VaultMigrationData.NoMigrationRequired, awaitItem())
+
+                vaultMigrationManager.verifyAndUpdateMigrationState(cipherList)
+
+                // Should still be NoMigrationRequired since no personal ciphers
+                expectNoEvents()
+            }
+        }
+
+    @Test
+    @Suppress("MaxLineLength")
+    fun `verifyAndUpdateMigrationState should emit NoMigrationRequired when cipher list is empty`() =
+        runTest {
+            fakeAuthDiskSource.userState = MOCK_USER_STATE
+
+            val mockPolicy = createMockPolicy(number = 1, type = PolicyTypeJson.PERSONAL_OWNERSHIP)
+            every {
+                policyManager.getActivePolicies(PolicyTypeJson.PERSONAL_OWNERSHIP)
+            } returns listOf(mockPolicy)
+
+            val cipherList = emptyList<SyncResponseJson.Cipher>()
+
+            vaultMigrationManager.vaultMigrationDataStateFlow.test {
+                assertEquals(VaultMigrationData.NoMigrationRequired, awaitItem())
+
+                vaultMigrationManager.verifyAndUpdateMigrationState(cipherList)
+
+                // Should still be NoMigrationRequired since cipher list is empty
+                expectNoEvents()
+            }
+        }
+
+    @Test
+    @Suppress("MaxLineLength")
+    fun `verifyAndUpdateMigrationState should emit NoMigrationRequired when organization ID is null`() =
+        runTest {
+            fakeAuthDiskSource.userState = MOCK_USER_STATE
+
+            val mockPolicy = createMockPolicy(number = 1, type = PolicyTypeJson.PERSONAL_OWNERSHIP)
+            every {
+                policyManager.getActivePolicies(PolicyTypeJson.PERSONAL_OWNERSHIP)
+            } returns listOf(mockPolicy)
+            every {
+                policyManager.getPersonalOwnershipPolicyOrganizationId()
+            } returns null
+
+            val personalCipher = createMockCipher(number = 1).copy(organizationId = null)
+            val cipherList = listOf(personalCipher)
+
+            vaultMigrationManager.vaultMigrationDataStateFlow.test {
+                assertEquals(VaultMigrationData.NoMigrationRequired, awaitItem())
+
+                vaultMigrationManager.verifyAndUpdateMigrationState(cipherList)
+
+                // Should still be NoMigrationRequired since organization ID is null
+                expectNoEvents()
+            }
+        }
+
+    @Test
+    @Suppress("MaxLineLength")
+    fun `verifyAndUpdateMigrationState should emit NoMigrationRequired when organizations list is null`() =
+        runTest {
+            val userId = "mockId-1"
+            fakeAuthDiskSource.userState = MOCK_USER_STATE
+
+            val mockPolicy = createMockPolicy(number = 1, type = PolicyTypeJson.PERSONAL_OWNERSHIP)
+            every {
+                policyManager.getActivePolicies(PolicyTypeJson.PERSONAL_OWNERSHIP)
+            } returns listOf(mockPolicy)
+            every {
+                policyManager.getPersonalOwnershipPolicyOrganizationId()
+            } returns "mockId-1"
+
+            // Store organizations as null in authDiskSource
+            val syncResponse = createMockSyncResponse(
+                number = 1,
+                profile = createMockProfile(number = 1, organizations = null),
+            )
+            fakeAuthDiskSource.storeOrganizations(userId = userId, organizations = null)
+
+            val personalCipher = createMockCipher(number = 1).copy(organizationId = null)
+            val cipherList = listOf(personalCipher)
+
+            vaultMigrationManager.vaultMigrationDataStateFlow.test {
+                assertEquals(VaultMigrationData.NoMigrationRequired, awaitItem())
+
+                vaultMigrationManager.verifyAndUpdateMigrationState(cipherList)
+
+                // Should still be NoMigrationRequired since organizations list is null
+                expectNoEvents()
+            }
+        }
+
+    @Test
+    @Suppress("MaxLineLength")
+    fun `verifyAndUpdateMigrationState should emit NoMigrationRequired when organization not found`() =
+        runTest {
+            val userId = "mockId-1"
+            fakeAuthDiskSource.userState = MOCK_USER_STATE
+
+            val mockPolicy = createMockPolicy(number = 1, type = PolicyTypeJson.PERSONAL_OWNERSHIP)
+            every {
+                policyManager.getActivePolicies(PolicyTypeJson.PERSONAL_OWNERSHIP)
+            } returns listOf(mockPolicy)
+            // Return a different org ID that doesn't match the one in the sync response
+            every {
+                policyManager.getPersonalOwnershipPolicyOrganizationId()
+            } returns "non-existent-org-id"
+
+            // Store organizations with a different ID
+            val syncResponse = createMockSyncResponse(number = 1)
+            fakeAuthDiskSource.storeOrganizations(
+                userId = userId,
+                organizations = syncResponse.profile.organizations,
+            )
+
+            val personalCipher = createMockCipher(number = 1).copy(organizationId = null)
+            val cipherList = listOf(personalCipher)
+
+            vaultMigrationManager.vaultMigrationDataStateFlow.test {
+                assertEquals(VaultMigrationData.NoMigrationRequired, awaitItem())
+
+                vaultMigrationManager.verifyAndUpdateMigrationState(cipherList)
+
+                // Should still be NoMigrationRequired since organization is not found
+                expectNoEvents()
+            }
+        }
+
+    @Test
+    fun `shouldMigrateVault returns true when all conditions are met`() {
+        fakeAuthDiskSource.userState = MOCK_USER_STATE
+
+        val mockPolicy = createMockPolicy(number = 1, type = PolicyTypeJson.PERSONAL_OWNERSHIP)
+        every {
+            policyManager.getActivePolicies(PolicyTypeJson.PERSONAL_OWNERSHIP)
+        } returns listOf(mockPolicy)
+
+        val result = vaultMigrationManager.shouldMigrateVault { true }
+
+        assertEquals(true, result)
+    }
+
+    @Test
+    fun `shouldMigrateVault returns false when no policies`() {
+        fakeAuthDiskSource.userState = MOCK_USER_STATE
+
+        every {
+            policyManager.getActivePolicies(PolicyTypeJson.PERSONAL_OWNERSHIP)
+        } returns emptyList()
+
+        val result = vaultMigrationManager.shouldMigrateVault { true }
+
+        assertEquals(false, result)
+    }
+
+    @Test
+    fun `shouldMigrateVault returns false when feature flag is disabled`() {
+        fakeAuthDiskSource.userState = MOCK_USER_STATE
+
+        val mockPolicy = createMockPolicy(number = 1, type = PolicyTypeJson.PERSONAL_OWNERSHIP)
+        every {
+            policyManager.getActivePolicies(PolicyTypeJson.PERSONAL_OWNERSHIP)
+        } returns listOf(mockPolicy)
+        every {
+            mockFeatureFlagManager.getFeatureFlag(FlagKey.MigrateMyVaultToMyItems)
+        } returns false
+
+        val result = vaultMigrationManager.shouldMigrateVault { true }
+
+        assertEquals(false, result)
+    }
+
+    @Test
+    fun `shouldMigrateVault returns false when no network`() {
+        fakeAuthDiskSource.userState = MOCK_USER_STATE
+
+        val mockPolicy = createMockPolicy(number = 1, type = PolicyTypeJson.PERSONAL_OWNERSHIP)
+        every {
+            policyManager.getActivePolicies(PolicyTypeJson.PERSONAL_OWNERSHIP)
+        } returns listOf(mockPolicy)
+        every { connectionManager.isNetworkConnected } returns false
+
+        val result = vaultMigrationManager.shouldMigrateVault { true }
+
+        assertEquals(false, result)
+    }
+
+    @Test
+    fun `shouldMigrateVault returns false when no personal items`() {
+        fakeAuthDiskSource.userState = MOCK_USER_STATE
+
+        val mockPolicy = createMockPolicy(number = 1, type = PolicyTypeJson.PERSONAL_OWNERSHIP)
+        every {
+            policyManager.getActivePolicies(PolicyTypeJson.PERSONAL_OWNERSHIP)
+        } returns listOf(mockPolicy)
+
+        val result = vaultMigrationManager.shouldMigrateVault { false }
+
+        assertEquals(false, result)
+    }
+}
+
+private val MOCK_USER_STATE = UserStateJson(
+    activeUserId = "mockId-1",
+    accounts = mapOf(
+        "mockId-1" to AccountJson(
+            profile = AccountJson.Profile(
+                userId = "mockId-1",
+                email = "test@bitwarden.com",
+                isEmailVerified = true,
+                name = "Bitwarden Tester",
+                stamp = null,
+                organizationId = null,
+                avatarColorHex = null,
+                hasPremium = false,
+                forcePasswordResetReason = null,
+                kdfType = null,
+                kdfIterations = null,
+                kdfMemory = null,
+                kdfParallelism = null,
+                userDecryptionOptions = null,
+                isTwoFactorEnabled = false,
+                creationDate = ZonedDateTime.parse("2024-09-13T01:00:00.00Z"),
+            ),
+            tokens = AccountTokensJson(
+                accessToken = "accessToken",
+                refreshToken = "refreshToken",
+            ),
+            settings = AccountJson.Settings(
+                environmentUrlData = null,
+            ),
+        ),
+    ),
+)
@@ -5,13 +5,11 @@ import app.cash.turbine.turbineScope
 import com.bitwarden.collections.CollectionView
 import com.bitwarden.core.InitOrgCryptoRequest
 import com.bitwarden.core.data.manager.dispatcher.FakeDispatcherManager
-import com.bitwarden.core.data.manager.model.FlagKey
 import com.bitwarden.core.data.repository.error.MissingPropertyException
 import com.bitwarden.core.data.repository.model.DataState
 import com.bitwarden.core.data.repository.util.bufferedMutableSharedFlow
 import com.bitwarden.core.data.util.asFailure
 import com.bitwarden.core.data.util.asSuccess
-import com.bitwarden.network.model.PolicyTypeJson
 import com.bitwarden.network.model.SyncResponseJson
 import com.bitwarden.network.model.createMockCipher
 import com.bitwarden.network.model.createMockCollection
@@ -37,10 +35,7 @@ import com.x8bit.bitwarden.data.auth.repository.model.LogoutReason
 import com.x8bit.bitwarden.data.platform.datasource.disk.SettingsDiskSource
 import com.x8bit.bitwarden.data.platform.error.NoActiveUserException
 import com.x8bit.bitwarden.data.platform.manager.DatabaseSchemeManager
-import com.x8bit.bitwarden.data.platform.manager.FeatureFlagManager
-import com.x8bit.bitwarden.data.platform.manager.PolicyManager
 import com.x8bit.bitwarden.data.platform.manager.PushManager
-import com.x8bit.bitwarden.data.platform.manager.network.NetworkConnectionManager
 import com.x8bit.bitwarden.data.vault.datasource.disk.VaultDiskSource
 import com.x8bit.bitwarden.data.vault.datasource.sdk.VaultSdkSource
 import com.x8bit.bitwarden.data.vault.datasource.sdk.model.InitializeCryptoResult
@@ -146,14 +141,12 @@ class VaultSyncManagerTest {
    private val databaseSchemeManager: DatabaseSchemeManager = mockk {
        every { databaseSchemeChangeFlow } returns mutableDatabaseSchemeChangeFlow
    }
-    private val policyManager: PolicyManager = mockk<PolicyManager> {
-        every { getActivePolicies(any()) } returns emptyList()
-    }
-    private val mockFeatureFlagManager = mockk<FeatureFlagManager>(relaxed = true) {
-        every { getFeatureFlag(FlagKey.MigrateMyVaultToMyItems) } returns true
-    }
-    private val connectionManager: NetworkConnectionManager = mockk<NetworkConnectionManager> {
-        every { isNetworkConnected } returns true
+    private val mutableVaultMigrationDataStateFlow =
+        MutableStateFlow<VaultMigrationData>(VaultMigrationData.NoMigrationRequired)
+    private val vaultMigrationManager: VaultMigrationManager = mockk {
+        every { vaultMigrationDataStateFlow } returns mutableVaultMigrationDataStateFlow
+        every { verifyAndUpdateMigrationState(any()) } just runs
+        every { shouldMigrateVault(any()) } returns false
    }

    private val vaultSyncManager: VaultSyncManager = VaultSyncManagerImpl(
@@ -165,9 +158,7 @@ class VaultSyncManagerTest {
        userLogoutManager = userLogoutManager,
        userStateManager = userStateManager,
        vaultLockManager = vaultLockManager,
-        policyManager = policyManager,
-        featureFlagManager = mockFeatureFlagManager,
-        connectionManager = connectionManager,
+        vaultMigrationManager = vaultMigrationManager,
        clock = clock,
        databaseSchemeManager = databaseSchemeManager,
        pushManager = pushManager,
@@ -1267,390 +1258,6 @@ class VaultSyncManagerTest {
            coVerify(exactly = 0) { syncService.sync() }
        }

-    @Test
-    fun `shouldMigratePersonalVaultFlow should initially emit false`() = runTest {
-        assertEquals(
-            VaultMigrationData.NoMigrationRequired,
-            vaultSyncManager.shouldMigratePersonalVaultFlow.value,
-        )
-    }
-
-    @Test
-    fun `shouldMigratePersonalVaultFlow should emit true after sync when all conditions are met`() =
-        runTest {
-            val userId = "mockId-1"
-            fakeAuthDiskSource.userState = MOCK_USER_STATE
-            setVaultToUnlocked(userId)
-
-            // Setup conditions for migration
-            val mockPolicy = createMockPolicy(number = 1, type = PolicyTypeJson.PERSONAL_OWNERSHIP)
-            every {
-                policyManager.getActivePolicies(PolicyTypeJson.PERSONAL_OWNERSHIP)
-            } returns listOf(mockPolicy)
-            every {
-                policyManager.getPersonalOwnershipPolicyOrganizationId()
-            } returns "mockId-1"
-            coEvery {
-                vaultSdkSource.initializeOrganizationCrypto(
-                    userId = userId,
-                    request = InitOrgCryptoRequest(
-                        organizationKeys = createMockOrganizationKeys(number = 1),
-                    ),
-                )
-            } returns InitializeCryptoResult.Success.asSuccess()
-
-            // Setup sync with personal ciphers (organizationId = null)
-            val personalCipher = createMockCipher(number = 1).copy(organizationId = null)
-            val syncResponse =
-                createMockSyncResponse(number = 1).copy(ciphers = listOf(personalCipher))
-            coEvery { syncService.sync() } returns syncResponse.asSuccess()
-            coEvery { vaultDiskSource.replaceVaultData(userId = userId, vault = any()) } just runs
-
-            setupEmptyDecryptionResults(userId)
-            setupVaultDiskSourceFlows(userId)
-
-            vaultSyncManager.shouldMigratePersonalVaultFlow.test {
-                assertEquals(VaultMigrationData.NoMigrationRequired, awaitItem())
-
-                vaultSyncManager.sync()
-
-                assertEquals(
-                    VaultMigrationData.MigrationRequired(
-                        organizationId = "mockId-1",
-                        organizationName = "mockName-1",
-                    ),
-                    awaitItem(),
-                )
-            }
-        }
-
-    @Test
-    fun `shouldMigratePersonalVaultFlow should emit false when no personal ownership policy`() =
-        runTest {
-            val userId = "mockId-1"
-            fakeAuthDiskSource.userState = MOCK_USER_STATE
-            setVaultToUnlocked(userId)
-
-            every {
-                policyManager.getActivePolicies(PolicyTypeJson.PERSONAL_OWNERSHIP)
-            } returns emptyList()
-
-            val personalCipher = createMockCipher(number = 1).copy(organizationId = null)
-            val syncResponse =
-                createMockSyncResponse(number = 1).copy(ciphers = listOf(personalCipher))
-            coEvery { syncService.sync() } returns syncResponse.asSuccess()
-            coEvery { vaultDiskSource.replaceVaultData(userId = userId, vault = any()) } just runs
-            coEvery {
-                vaultSdkSource.initializeOrganizationCrypto(
-                    userId = userId,
-                    request = InitOrgCryptoRequest(
-                        organizationKeys = createMockOrganizationKeys(number = 1),
-                    ),
-                )
-            } returns InitializeCryptoResult.Success.asSuccess()
-
-            setupEmptyDecryptionResults(userId)
-            setupVaultDiskSourceFlows(userId)
-
-            vaultSyncManager.shouldMigratePersonalVaultFlow.test {
-                assertEquals(VaultMigrationData.NoMigrationRequired, awaitItem())
-
-                vaultSyncManager.sync()
-
-                // Should still be false since no policy
-                expectNoEvents()
-            }
-        }
-
-    @Test
-    fun `shouldMigratePersonalVaultFlow should emit false when feature flag is disabled`() =
-        runTest {
-            val userId = "mockId-1"
-            fakeAuthDiskSource.userState = MOCK_USER_STATE
-            setVaultToUnlocked(userId)
-
-            val mockPolicy = createMockPolicy(number = 1, type = PolicyTypeJson.PERSONAL_OWNERSHIP)
-            every {
-                policyManager.getActivePolicies(PolicyTypeJson.PERSONAL_OWNERSHIP)
-            } returns listOf(mockPolicy)
-
-            every {
-                mockFeatureFlagManager.getFeatureFlag(FlagKey.MigrateMyVaultToMyItems)
-            } returns false
-
-            val personalCipher = createMockCipher(number = 1).copy(organizationId = null)
-            val syncResponse =
-                createMockSyncResponse(number = 1).copy(ciphers = listOf(personalCipher))
-            coEvery { syncService.sync() } returns syncResponse.asSuccess()
-            coEvery { vaultDiskSource.replaceVaultData(userId = userId, vault = any()) } just runs
-            coEvery {
-                vaultSdkSource.initializeOrganizationCrypto(
-                    userId = userId,
-                    request = InitOrgCryptoRequest(
-                        organizationKeys = createMockOrganizationKeys(number = 1),
-                    ),
-                )
-            } returns InitializeCryptoResult.Success.asSuccess()
-
-            setupEmptyDecryptionResults(userId)
-            setupVaultDiskSourceFlows(userId)
-
-            vaultSyncManager.shouldMigratePersonalVaultFlow.test {
-                assertEquals(VaultMigrationData.NoMigrationRequired, awaitItem())
-
-                vaultSyncManager.sync()
-
-                // Should still be false since feature flag is disabled
-                expectNoEvents()
-            }
-        }
-
-    @Test
-    fun `shouldMigratePersonalVaultFlow should emit false when no network connection`() =
-        runTest {
-            val userId = "mockId-1"
-            fakeAuthDiskSource.userState = MOCK_USER_STATE
-            setVaultToUnlocked(userId)
-
-            val mockPolicy = createMockPolicy(number = 1, type = PolicyTypeJson.PERSONAL_OWNERSHIP)
-            every {
-                policyManager.getActivePolicies(PolicyTypeJson.PERSONAL_OWNERSHIP)
-            } returns listOf(mockPolicy)
-
-            every { connectionManager.isNetworkConnected } returns false
-
-            val personalCipher = createMockCipher(number = 1).copy(organizationId = null)
-            val syncResponse =
-                createMockSyncResponse(number = 1).copy(ciphers = listOf(personalCipher))
-            coEvery { syncService.sync() } returns syncResponse.asSuccess()
-            coEvery { vaultDiskSource.replaceVaultData(userId = userId, vault = any()) } just runs
-            coEvery {
-                vaultSdkSource.initializeOrganizationCrypto(
-                    userId = userId,
-                    request = InitOrgCryptoRequest(
-                        organizationKeys = createMockOrganizationKeys(number = 1),
-                    ),
-                )
-            } returns InitializeCryptoResult.Success.asSuccess()
-
-            setupEmptyDecryptionResults(userId)
-            setupVaultDiskSourceFlows(userId)
-
-            vaultSyncManager.shouldMigratePersonalVaultFlow.test {
-                assertEquals(VaultMigrationData.NoMigrationRequired, awaitItem())
-
-                vaultSyncManager.sync()
-
-                // Should still be false since no network
-                expectNoEvents()
-            }
-        }
-
-    @Test
-    fun `shouldMigratePersonalVaultFlow should emit false when no personal ciphers exist`() =
-        runTest {
-            val userId = "mockId-1"
-            fakeAuthDiskSource.userState = MOCK_USER_STATE
-            setVaultToUnlocked(userId)
-
-            val mockPolicy = createMockPolicy(number = 1, type = PolicyTypeJson.PERSONAL_OWNERSHIP)
-            every {
-                policyManager.getActivePolicies(PolicyTypeJson.PERSONAL_OWNERSHIP)
-            } returns listOf(mockPolicy)
-
-            // All ciphers belong to organizations
-            val orgCipher = createMockCipher(number = 1).copy(organizationId = "org-id")
-            val syncResponse = createMockSyncResponse(number = 1).copy(ciphers = listOf(orgCipher))
-            coEvery { syncService.sync() } returns syncResponse.asSuccess()
-            coEvery { vaultDiskSource.replaceVaultData(userId = userId, vault = any()) } just runs
-            coEvery {
-                vaultSdkSource.initializeOrganizationCrypto(
-                    userId = userId,
-                    request = InitOrgCryptoRequest(
-                        organizationKeys = createMockOrganizationKeys(number = 1),
-                    ),
-                )
-            } returns InitializeCryptoResult.Success.asSuccess()
-
-            setupEmptyDecryptionResults(userId)
-            setupVaultDiskSourceFlows(userId)
-
-            vaultSyncManager.shouldMigratePersonalVaultFlow.test {
-                assertEquals(VaultMigrationData.NoMigrationRequired, awaitItem())
-
-                vaultSyncManager.sync()
-
-                // Should still be false since no personal ciphers
-                expectNoEvents()
-            }
-        }
-
-    @Test
-    fun `shouldMigratePersonalVaultFlow should emit false when cipher list is null`() =
-        runTest {
-            val userId = "mockId-1"
-            fakeAuthDiskSource.userState = MOCK_USER_STATE
-            setVaultToUnlocked(userId)
-
-            val mockPolicy = createMockPolicy(number = 1, type = PolicyTypeJson.PERSONAL_OWNERSHIP)
-            every {
-                policyManager.getActivePolicies(PolicyTypeJson.PERSONAL_OWNERSHIP)
-            } returns listOf(mockPolicy)
-
-            val syncResponse = createMockSyncResponse(number = 1).copy(ciphers = null)
-            coEvery { syncService.sync() } returns syncResponse.asSuccess()
-            coEvery { vaultDiskSource.replaceVaultData(userId = userId, vault = any()) } just runs
-
-            coEvery {
-                vaultSdkSource.initializeOrganizationCrypto(
-                    userId = userId,
-                    request = InitOrgCryptoRequest(
-                        organizationKeys = createMockOrganizationKeys(number = 1),
-                    ),
-                )
-            } returns InitializeCryptoResult.Success.asSuccess()
-
-            setupEmptyDecryptionResults(userId)
-            setupVaultDiskSourceFlows(userId)
-
-            vaultSyncManager.shouldMigratePersonalVaultFlow.test {
-                assertEquals(VaultMigrationData.NoMigrationRequired, awaitItem())
-
-                vaultSyncManager.sync()
-
-                // Should still be false since ciphers is null
-                expectNoEvents()
-            }
-        }
-
-    @Test
-    fun `shouldMigratePersonalVaultFlow should emit false when organization ID is null`() =
-        runTest {
-            val userId = "mockId-1"
-            fakeAuthDiskSource.userState = MOCK_USER_STATE
-            setVaultToUnlocked(userId)
-
-            val mockPolicy = createMockPolicy(number = 1, type = PolicyTypeJson.PERSONAL_OWNERSHIP)
-            every {
-                policyManager.getActivePolicies(PolicyTypeJson.PERSONAL_OWNERSHIP)
-            } returns listOf(mockPolicy)
-            every {
-                policyManager.getPersonalOwnershipPolicyOrganizationId()
-            } returns null
-
-            val personalCipher = createMockCipher(number = 1).copy(organizationId = null)
-            val syncResponse =
-                createMockSyncResponse(number = 1).copy(ciphers = listOf(personalCipher))
-            coEvery { syncService.sync() } returns syncResponse.asSuccess()
-            coEvery { vaultDiskSource.replaceVaultData(userId = userId, vault = any()) } just runs
-            coEvery {
-                vaultSdkSource.initializeOrganizationCrypto(
-                    userId = userId,
-                    request = InitOrgCryptoRequest(
-                        organizationKeys = createMockOrganizationKeys(number = 1),
-                    ),
-                )
-            } returns InitializeCryptoResult.Success.asSuccess()
-
-            setupEmptyDecryptionResults(userId)
-            setupVaultDiskSourceFlows(userId)
-
-            vaultSyncManager.shouldMigratePersonalVaultFlow.test {
-                assertEquals(VaultMigrationData.NoMigrationRequired, awaitItem())
-
-                vaultSyncManager.sync()
-
-                // Should still be NoMigrationRequired since organization ID is null
-                expectNoEvents()
-            }
-        }
-
-    @Test
-    fun `shouldMigratePersonalVaultFlow should emit false when organizations list is null`() =
-        runTest {
-            val userId = "mockId-1"
-            fakeAuthDiskSource.userState = MOCK_USER_STATE
-            setVaultToUnlocked(userId)
-
-            val mockPolicy = createMockPolicy(number = 1, type = PolicyTypeJson.PERSONAL_OWNERSHIP)
-            every {
-                policyManager.getActivePolicies(PolicyTypeJson.PERSONAL_OWNERSHIP)
-            } returns listOf(mockPolicy)
-            every {
-                policyManager.getPersonalOwnershipPolicyOrganizationId()
-            } returns "mockId-1"
-
-            val personalCipher = createMockCipher(number = 1).copy(organizationId = null)
-            val syncResponse = createMockSyncResponse(
-                number = 1,
-                profile = createMockProfile(number = 1, organizations = null),
-            ).copy(ciphers = listOf(personalCipher))
-            coEvery { syncService.sync() } returns syncResponse.asSuccess()
-            coEvery { vaultDiskSource.replaceVaultData(userId = userId, vault = any()) } just runs
-            coEvery {
-                vaultSdkSource.initializeOrganizationCrypto(
-                    userId = userId,
-                    request = InitOrgCryptoRequest(
-                        organizationKeys = emptyMap(),
-                    ),
-                )
-            } returns InitializeCryptoResult.Success.asSuccess()
-
-            setupEmptyDecryptionResults(userId)
-            setupVaultDiskSourceFlows(userId)
-
-            vaultSyncManager.shouldMigratePersonalVaultFlow.test {
-                assertEquals(VaultMigrationData.NoMigrationRequired, awaitItem())
-
-                vaultSyncManager.sync()
-
-                // Should still be NoMigrationRequired since organizations list is null
-                expectNoEvents()
-            }
-        }
-
-    @Test
-    fun `shouldMigratePersonalVaultFlow should emit false when organization not found`() =
-        runTest {
-            val userId = "mockId-1"
-            fakeAuthDiskSource.userState = MOCK_USER_STATE
-            setVaultToUnlocked(userId)
-
-            val mockPolicy = createMockPolicy(number = 1, type = PolicyTypeJson.PERSONAL_OWNERSHIP)
-            every {
-                policyManager.getActivePolicies(PolicyTypeJson.PERSONAL_OWNERSHIP)
-            } returns listOf(mockPolicy)
-            every {
-                policyManager.getPersonalOwnershipPolicyOrganizationId()
-            } returns "non-existent-org-id"
-
-            val personalCipher = createMockCipher(number = 1).copy(organizationId = null)
-            val syncResponse =
-                createMockSyncResponse(number = 1).copy(ciphers = listOf(personalCipher))
-            coEvery { syncService.sync() } returns syncResponse.asSuccess()
-            coEvery { vaultDiskSource.replaceVaultData(userId = userId, vault = any()) } just runs
-            coEvery {
-                vaultSdkSource.initializeOrganizationCrypto(
-                    userId = userId,
-                    request = InitOrgCryptoRequest(
-                        organizationKeys = createMockOrganizationKeys(number = 1),
-                    ),
-                )
-            } returns InitializeCryptoResult.Success.asSuccess()
-
-            setupEmptyDecryptionResults(userId)
-            setupVaultDiskSourceFlows(userId)
-
-            vaultSyncManager.shouldMigratePersonalVaultFlow.test {
-                assertEquals(VaultMigrationData.NoMigrationRequired, awaitItem())
-
-                vaultSyncManager.sync()
-
-                // Should still be NoMigrationRequired since organization is not found
-                expectNoEvents()
-            }
-        }
-
    //region Helper functions

    /**
@@ -28,8 +28,8 @@ import com.x8bit.bitwarden.data.platform.manager.SpecialCircumstanceManagerImpl
 import com.x8bit.bitwarden.data.platform.manager.model.CompleteRegistrationData
 import com.x8bit.bitwarden.data.platform.manager.model.FirstTimeState
 import com.x8bit.bitwarden.data.platform.manager.model.SpecialCircumstance
+import com.x8bit.bitwarden.data.vault.manager.VaultMigrationManager
 import com.x8bit.bitwarden.data.vault.manager.model.VaultMigrationData
-import com.x8bit.bitwarden.data.vault.repository.VaultRepository
 import com.x8bit.bitwarden.ui.tools.feature.send.model.SendItemType
 import io.mockk.every
 import io.mockk.mockk
@@ -63,10 +63,10 @@ class RootNavViewModelTest : BaseViewModelTest() {
            dispatcherManager = FakeDispatcherManager(),
        )

-    private val mutableShouldMigratePersonalVaultFlow =
+    private val mutableVaultMigrationDataStateFlow =
        MutableStateFlow<VaultMigrationData>(VaultMigrationData.NoMigrationRequired)
-    private val vaultRepository = mockk<VaultRepository> {
-        every { shouldMigratePersonalVaultFlow } returns mutableShouldMigratePersonalVaultFlow
+    private val vaultMigrationManager = mockk<VaultMigrationManager> {
+        every { vaultMigrationDataStateFlow } returns mutableVaultMigrationDataStateFlow
    }

    @BeforeEach
@@ -1589,8 +1589,8 @@ class RootNavViewModelTest : BaseViewModelTest() {

    @Suppress("MaxLineLength")
    @Test
-    fun `when shouldMigratePersonalVaultFlow emits true the nav state should be MigrateToMyItems`() {
-        mutableShouldMigratePersonalVaultFlow.value = MOCK_VAULT_MIGRATION_DATA
+    fun `when vaultMigrationDataStateFlow emits true the nav state should be MigrateToMyItems`() {
+        mutableVaultMigrationDataStateFlow.value = MOCK_VAULT_MIGRATION_DATA
        mutableUserStateFlow.tryEmit(MOCK_VAULT_UNLOCKED_USER_STATE)
        val viewModel = createViewModel()

@@ -1605,7 +1605,7 @@ class RootNavViewModelTest : BaseViewModelTest() {

    @Test
    fun `when vault is locked MigrateToMyItems should not show even if flow emits true`() {
-        mutableShouldMigratePersonalVaultFlow.value = MOCK_VAULT_MIGRATION_DATA
+        mutableVaultMigrationDataStateFlow.value = MOCK_VAULT_MIGRATION_DATA

        // Vault is locked
        mutableUserStateFlow.tryEmit(
@@ -1624,8 +1624,8 @@ class RootNavViewModelTest : BaseViewModelTest() {
    }

    @Test
-    fun `when shouldMigratePersonalVaultFlow emits false MigrateToMyItems should not show`() {
-        mutableShouldMigratePersonalVaultFlow.value = VaultMigrationData.NoMigrationRequired
+    fun `when vaultMigrationDataStateFlow emits false MigrateToMyItems should not show`() {
+        mutableVaultMigrationDataStateFlow.value = VaultMigrationData.NoMigrationRequired
        mutableUserStateFlow.tryEmit(MOCK_VAULT_UNLOCKED_USER_STATE)
        val viewModel = createViewModel()

@@ -1639,7 +1639,7 @@ class RootNavViewModelTest : BaseViewModelTest() {
        RootNavViewModel(
            authRepository = authRepository,
            specialCircumstanceManager = specialCircumstanceManager,
-            vaultRepository = vaultRepository,
+            vaultMigrationManager = vaultMigrationManager,
        )
 }