PM-42381: chore: Separate AuthState and organizations from AuthRepository (#7296)

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/manager/AuthStateManager.kt

@@ -0,0 +1,14 @@
+package com.x8bit.bitwarden.data.auth.manager
+
+import com.x8bit.bitwarden.data.auth.repository.model.AuthState
+import kotlinx.coroutines.flow.StateFlow
+
+/**
+ * Manages the authentication state.
+ */
+interface AuthStateManager {
+    /**
+     * Models the current auth state.
+     */
+    val authStateFlow: StateFlow<AuthState>
+}

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/manager/AuthStateManagerImpl.kt

@@ -0,0 +1,47 @@
+package com.x8bit.bitwarden.data.auth.manager
+
+import com.bitwarden.core.data.manager.dispatcher.DispatcherManager
+import com.x8bit.bitwarden.data.auth.datasource.disk.AuthDiskSource
+import com.x8bit.bitwarden.data.auth.repository.model.AuthState
+import com.x8bit.bitwarden.data.auth.repository.util.activeUserIdChangesFlow
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.SharingStarted
+import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.flatMapLatest
+import kotlinx.coroutines.flow.flowOf
+import kotlinx.coroutines.flow.map
+import kotlinx.coroutines.flow.stateIn
+
+/**
+ * The default implementation of the [AuthStateManager].
+ */
+internal class AuthStateManagerImpl(
+    private val authDiskSource: AuthDiskSource,
+    dispatcherManager: DispatcherManager,
+) : AuthStateManager {
+    private val unconfinedScope = CoroutineScope(context = dispatcherManager.unconfined)
+
+    @OptIn(ExperimentalCoroutinesApi::class)
+    override val authStateFlow: StateFlow<AuthState> = authDiskSource
+        .activeUserIdChangesFlow
+        .flatMapLatest { activeUserId ->
+            activeUserId
+                ?.let { userId ->
+                    authDiskSource
+                        .getAccountTokensFlow(userId = userId)
+                        .map { accountTokens ->
+                            accountTokens
+                                ?.accessToken
+                                ?.let { AuthState.Authenticated(accessToken = it) }
+                                ?: AuthState.Unauthenticated
+                        }
+                }
+                ?: flowOf(AuthState.Unauthenticated)
+        }
+        .stateIn(
+            scope = unconfinedScope,
+            started = SharingStarted.Eagerly,
+            initialValue = AuthState.Uninitialized,
+        )
+}

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/manager/OrganizationManager.kt

@@ -0,0 +1,25 @@
+package com.x8bit.bitwarden.data.auth.manager
+
+import com.x8bit.bitwarden.data.auth.repository.model.LeaveOrganizationResult
+import com.x8bit.bitwarden.data.auth.repository.model.Organization
+import com.x8bit.bitwarden.data.auth.repository.model.RevokeFromOrganizationResult
+
+/**
+ * Manager used to manage organizations.
+ */
+interface OrganizationManager {
+    /**
+     * The organization for the active user.
+     */
+    val organizations: List<Organization>
+
+    /**
+     * Leaves the organization that matches the given [organizationId]
+     */
+    suspend fun leaveOrganization(organizationId: String): LeaveOrganizationResult
+
+    /**
+     * Revokes self from the organization that matches the given [organizationId]
+     */
+    suspend fun revokeFromOrganization(organizationId: String): RevokeFromOrganizationResult
+}

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/manager/OrganizationManagerImpl.kt

@@ -0,0 +1,42 @@
+package com.x8bit.bitwarden.data.auth.manager
+
+import com.bitwarden.network.model.OrganizationStatusType
+import com.bitwarden.network.service.OrganizationService
+import com.x8bit.bitwarden.data.auth.datasource.disk.AuthDiskSource
+import com.x8bit.bitwarden.data.auth.repository.model.LeaveOrganizationResult
+import com.x8bit.bitwarden.data.auth.repository.model.Organization
+import com.x8bit.bitwarden.data.auth.repository.model.RevokeFromOrganizationResult
+import com.x8bit.bitwarden.data.auth.repository.util.toOrganizations
+
+/**
+ * The default implementation of the [OrganizationManager].
+ */
+internal class OrganizationManagerImpl(
+    private val authDiskSource: AuthDiskSource,
+    private val organizationService: OrganizationService,
+) : OrganizationManager {
+    private val activeUserId: String? get() = authDiskSource.userState?.activeUserId
+
+    override val organizations: List<Organization>
+        get() = activeUserId
+            ?.let { authDiskSource.getOrganizations(userId = it) }
+            ?.filter { it.status == OrganizationStatusType.CONFIRMED }
+            .orEmpty()
+            .toOrganizations()
+
+    override suspend fun leaveOrganization(
+        organizationId: String,
+    ): LeaveOrganizationResult =
+        organizationService.leaveOrganization(organizationId = organizationId).fold(
+            onSuccess = { LeaveOrganizationResult.Success },
+            onFailure = { LeaveOrganizationResult.Error(error = it) },
+        )
+
+    override suspend fun revokeFromOrganization(
+        organizationId: String,
+    ): RevokeFromOrganizationResult =
+        organizationService.revokeFromOrganization(organizationId = organizationId).fold(
+            onSuccess = { RevokeFromOrganizationResult.Success },
+            onFailure = { RevokeFromOrganizationResult.Error(error = it) },
+        )
+}

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/manager/di/AuthManagerModule.kt

@@ -7,6 +7,7 @@ import com.bitwarden.network.service.AccountsService
 import com.bitwarden.network.service.AuthRequestsService
 import com.bitwarden.network.service.DevicesService
 import com.bitwarden.network.service.NewAuthRequestService
+import com.bitwarden.network.service.OrganizationService
 import com.x8bit.bitwarden.data.auth.datasource.disk.AuthDiskSource
 import com.x8bit.bitwarden.data.auth.datasource.sdk.AuthSdkSource
 import com.x8bit.bitwarden.data.auth.manager.AddTotpItemFromAuthenticatorManager
@@ -15,25 +16,34 @@ import com.x8bit.bitwarden.data.auth.manager.AuthRequestManager
 import com.x8bit.bitwarden.data.auth.manager.AuthRequestManagerImpl
 import com.x8bit.bitwarden.data.auth.manager.AuthRequestNotificationManager
 import com.x8bit.bitwarden.data.auth.manager.AuthRequestNotificationManagerImpl
+import com.x8bit.bitwarden.data.auth.manager.AuthStateManager
+import com.x8bit.bitwarden.data.auth.manager.AuthStateManagerImpl
 import com.x8bit.bitwarden.data.auth.manager.AuthTokenManager
 import com.x8bit.bitwarden.data.auth.manager.AuthTokenManagerImpl
 import com.x8bit.bitwarden.data.auth.manager.KdfManager
 import com.x8bit.bitwarden.data.auth.manager.KdfManagerImpl
 import com.x8bit.bitwarden.data.auth.manager.KeyConnectorManager
 import com.x8bit.bitwarden.data.auth.manager.KeyConnectorManagerImpl
+import com.x8bit.bitwarden.data.auth.manager.OrganizationManager
+import com.x8bit.bitwarden.data.auth.manager.OrganizationManagerImpl
 import com.x8bit.bitwarden.data.auth.manager.TrustedDeviceManager
 import com.x8bit.bitwarden.data.auth.manager.TrustedDeviceManagerImpl
 import com.x8bit.bitwarden.data.auth.manager.UserLogoutManager
 import com.x8bit.bitwarden.data.auth.manager.UserLogoutManagerImpl
+import com.x8bit.bitwarden.data.auth.manager.UserStateManager
+import com.x8bit.bitwarden.data.auth.manager.UserStateManagerImpl
 import com.x8bit.bitwarden.data.platform.datasource.disk.PushDiskSource
 import com.x8bit.bitwarden.data.platform.datasource.disk.SettingsDiskSource
 import com.x8bit.bitwarden.data.platform.manager.CredentialExchangeRegistryManager
 import com.x8bit.bitwarden.data.platform.manager.FeatureFlagManager
+import com.x8bit.bitwarden.data.platform.manager.FirstTimeActionManager
+import com.x8bit.bitwarden.data.platform.manager.PolicyManager
 import com.x8bit.bitwarden.data.platform.manager.PushManager
 import com.x8bit.bitwarden.data.tools.generator.datasource.disk.GeneratorDiskSource
 import com.x8bit.bitwarden.data.tools.generator.datasource.disk.PasswordHistoryDiskSource
 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.VaultLockManager
 import dagger.Module
 import dagger.Provides
 import dagger.hilt.InstallIn
@@ -83,6 +93,16 @@ object AuthManagerModule {
             authDiskSource = authDiskSource,
         )
 
+    @Provides
+    @Singleton
+    fun providesAuthStateManager(
+        authDiskSource: AuthDiskSource,
+        dispatcherManager: DispatcherManager,
+    ): AuthStateManager = AuthStateManagerImpl(
+        authDiskSource = authDiskSource,
+        dispatcherManager = dispatcherManager,
+    )
+
     @Provides
     @Singleton
     fun provideKeyConnectorManager(
@@ -162,4 +182,30 @@ object AuthManagerModule {
         accountsService = accountsService,
         featureFlagManager = featureFlagManager,
     )
+
+    @Provides
+    @Singleton
+    fun providesOrganizationManager(
+        authDiskSource: AuthDiskSource,
+        organizationService: OrganizationService,
+    ): OrganizationManager = OrganizationManagerImpl(
+        authDiskSource = authDiskSource,
+        organizationService = organizationService,
+    )
+
+    @Provides
+    @Singleton
+    fun providesUserStateManager(
+        authDiskSource: AuthDiskSource,
+        firstTimeActionManager: FirstTimeActionManager,
+        vaultLockManager: VaultLockManager,
+        policyManager: PolicyManager,
+        dispatcherManager: DispatcherManager,
+    ): UserStateManager = UserStateManagerImpl(
+        authDiskSource = authDiskSource,
+        firstTimeActionManager = firstTimeActionManager,
+        vaultLockManager = vaultLockManager,
+        policyManager = policyManager,
+        dispatcherManager = dispatcherManager,
+    )
 }

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/repository/AuthRepository.kt

@@ -4,19 +4,18 @@ import com.bitwarden.network.model.GetTokenResponseJson
 import com.bitwarden.network.model.TwoFactorDataModel
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.OnboardingStatus
 import com.x8bit.bitwarden.data.auth.manager.AuthRequestManager
+import com.x8bit.bitwarden.data.auth.manager.AuthStateManager
 import com.x8bit.bitwarden.data.auth.manager.KdfManager
+import com.x8bit.bitwarden.data.auth.manager.OrganizationManager
 import com.x8bit.bitwarden.data.auth.manager.UserStateManager
-import com.x8bit.bitwarden.data.auth.repository.model.AuthState
 import com.x8bit.bitwarden.data.auth.repository.model.BreachCountResult
 import com.x8bit.bitwarden.data.auth.repository.model.DeleteAccountResult
 import com.x8bit.bitwarden.data.auth.repository.model.EmailTokenResult
 import com.x8bit.bitwarden.data.auth.repository.model.GetDevicesResult
 import com.x8bit.bitwarden.data.auth.repository.model.KnownDeviceResult
-import com.x8bit.bitwarden.data.auth.repository.model.LeaveOrganizationResult
 import com.x8bit.bitwarden.data.auth.repository.model.LoginResult
 import com.x8bit.bitwarden.data.auth.repository.model.LogoutReason
 import com.x8bit.bitwarden.data.auth.repository.model.NewSsoUserResult
-import com.x8bit.bitwarden.data.auth.repository.model.Organization
 import com.x8bit.bitwarden.data.auth.repository.model.PasswordHintResult
 import com.x8bit.bitwarden.data.auth.repository.model.PrevalidateSsoResult
 import com.x8bit.bitwarden.data.auth.repository.model.RegisterResult
@@ -24,7 +23,6 @@ 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
-import com.x8bit.bitwarden.data.auth.repository.model.RevokeFromOrganizationResult
 import com.x8bit.bitwarden.data.auth.repository.model.SendVerificationEmailResult
 import com.x8bit.bitwarden.data.auth.repository.model.SetPasswordResult
 import com.x8bit.bitwarden.data.auth.repository.model.SwitchAccountResult
@@ -41,7 +39,6 @@ import com.x8bit.bitwarden.data.platform.datasource.network.authenticator.Authen
 import com.x8bit.bitwarden.data.platform.manager.BiometricsEncryptionManager
 import com.x8bit.bitwarden.data.platform.manager.policy.PasswordPolicyManager
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.StateFlow
 
 /**
  * Provides an API for observing and modifying authentication state.
@@ -50,15 +47,12 @@ import kotlinx.coroutines.flow.StateFlow
 interface AuthRepository :
     AuthenticatorProvider,
     AuthRequestManager,
+    AuthStateManager,
     BiometricsEncryptionManager,
     KdfManager,
+    OrganizationManager,
     PasswordPolicyManager,
     UserStateManager {
-    /**
-     * Models the current auth state.
-     */
-    val authStateFlow: StateFlow<AuthState>
-
     /**
      * Flow of the current [DuoCallbackTokenResult]. Subscribers should listen to the flow
      * in order to receive updates whenever [setDuoCallbackTokenResult] is called.
@@ -120,11 +114,6 @@ interface AuthRepository :
      */
     var shouldTrustDevice: Boolean
 
-    /**
-     * The organization for the active user.
-     */
-    val organizations: List<Organization>
-
     /**
      * Whether the welcome carousel should be displayed, based on the feature flag and
      * whether the user has ever logged in or created an account before.
@@ -393,18 +382,4 @@ interface AuthRepository :
      * Update the value of the onboarding status for the user.
      */
     fun setOnboardingStatus(status: OnboardingStatus)
-
-    /**
-     * Leaves the organization that matches the given [organizationId]
-     */
-    suspend fun leaveOrganization(
-        organizationId: String,
-    ): LeaveOrganizationResult
-
-    /**
-     * Revokes self from the organization that matches the given [organizationId]
-     */
-    suspend fun revokeFromOrganization(
-        organizationId: String,
-    ): RevokeFromOrganizationResult
 }

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryImpl.kt

@@ -25,7 +25,6 @@ import com.bitwarden.network.model.GetTokenResponseJson
 import com.bitwarden.network.model.IdentityTokenAuthModel
 import com.bitwarden.network.model.OrganizationAutoEnrollStatusResponseJson
 import com.bitwarden.network.model.OrganizationKeysResponseJson
-import com.bitwarden.network.model.OrganizationStatusType
 import com.bitwarden.network.model.OrganizationType
 import com.bitwarden.network.model.PasswordHintResponseJson
 import com.bitwarden.network.model.PrevalidateSsoResponseJson
@@ -61,23 +60,22 @@ import com.x8bit.bitwarden.data.auth.datasource.network.model.DeviceDataModel
 import com.x8bit.bitwarden.data.auth.datasource.sdk.AuthSdkSource
 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.AuthStateManager
 import com.x8bit.bitwarden.data.auth.manager.KdfManager
 import com.x8bit.bitwarden.data.auth.manager.KeyConnectorManager
+import com.x8bit.bitwarden.data.auth.manager.OrganizationManager
 import com.x8bit.bitwarden.data.auth.manager.TrustedDeviceManager
 import com.x8bit.bitwarden.data.auth.manager.UserLogoutManager
 import com.x8bit.bitwarden.data.auth.manager.UserStateManager
 import com.x8bit.bitwarden.data.auth.manager.model.MigrateExistingUserToKeyConnectorResult
-import com.x8bit.bitwarden.data.auth.repository.model.AuthState
 import com.x8bit.bitwarden.data.auth.repository.model.BreachCountResult
 import com.x8bit.bitwarden.data.auth.repository.model.DeleteAccountResult
 import com.x8bit.bitwarden.data.auth.repository.model.EmailTokenResult
 import com.x8bit.bitwarden.data.auth.repository.model.GetDevicesResult
 import com.x8bit.bitwarden.data.auth.repository.model.KnownDeviceResult
-import com.x8bit.bitwarden.data.auth.repository.model.LeaveOrganizationResult
 import com.x8bit.bitwarden.data.auth.repository.model.LoginResult
 import com.x8bit.bitwarden.data.auth.repository.model.LogoutReason
 import com.x8bit.bitwarden.data.auth.repository.model.NewSsoUserResult
-import com.x8bit.bitwarden.data.auth.repository.model.Organization
 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
@@ -86,7 +84,6 @@ 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
-import com.x8bit.bitwarden.data.auth.repository.model.RevokeFromOrganizationResult
 import com.x8bit.bitwarden.data.auth.repository.model.SendVerificationEmailResult
 import com.x8bit.bitwarden.data.auth.repository.model.SetPasswordResult
 import com.x8bit.bitwarden.data.auth.repository.model.SwitchAccountResult
@@ -100,10 +97,8 @@ import com.x8bit.bitwarden.data.auth.repository.util.CookieCallbackResult
 import com.x8bit.bitwarden.data.auth.repository.util.DuoCallbackTokenResult
 import com.x8bit.bitwarden.data.auth.repository.util.SsoCallbackResult
 import com.x8bit.bitwarden.data.auth.repository.util.WebAuthResult
-import com.x8bit.bitwarden.data.auth.repository.util.activeUserIdChangesFlow
 import com.x8bit.bitwarden.data.auth.repository.util.toAccountCryptographicState
 import com.x8bit.bitwarden.data.auth.repository.util.toDeviceInfo
-import com.x8bit.bitwarden.data.auth.repository.util.toOrganizations
 import com.x8bit.bitwarden.data.auth.repository.util.toPolicyInformation
 import com.x8bit.bitwarden.data.auth.repository.util.toSdkParams
 import com.x8bit.bitwarden.data.auth.repository.util.toUserState
@@ -130,20 +125,13 @@ import com.x8bit.bitwarden.data.vault.repository.model.onVaultUnlockSuccess
 import com.x8bit.bitwarden.data.vault.repository.util.toSdkMasterPasswordUnlock
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.flow.Flow
-import kotlinx.coroutines.flow.SharingStarted
-import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asSharedFlow
 import kotlinx.coroutines.flow.combine
-import kotlinx.coroutines.flow.flatMapLatest
-import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.flow.launchIn
-import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.receiveAsFlow
-import kotlinx.coroutines.flow.stateIn
 import timber.log.Timber
 import java.time.Clock
 import javax.inject.Singleton
@@ -177,14 +165,18 @@ internal class AuthRepositoryImpl(
     private val kdfManager: KdfManager,
     private val toastManager: ToastManager,
     private val featureFlagManager: FeatureFlagManager,
+    authStateManager: AuthStateManager,
+    organizationManager: OrganizationManager,
     passwordPolicyManager: PasswordPolicyManager,
     logsManager: LogsManager,
     pushManager: PushManager,
     dispatcherManager: DispatcherManager,
 ) : AuthRepository,
     AuthRequestManager by authRequestManager,
+    AuthStateManager by authStateManager,
     BiometricsEncryptionManager by biometricsEncryptionManager,
     KdfManager by kdfManager,
+    OrganizationManager by organizationManager,
     PasswordPolicyManager by passwordPolicyManager,
     UserStateManager by userStateManager {
     /**
@@ -230,29 +222,6 @@ internal class AuthRepositoryImpl(
     override val ssoOrganizationIdentifier: String? get() = organizationIdentifier
     override val activeUserId: String? get() = authDiskSource.userState?.activeUserId
 
-    @OptIn(ExperimentalCoroutinesApi::class)
-    override val authStateFlow: StateFlow<AuthState> = authDiskSource
-        .activeUserIdChangesFlow
-        .flatMapLatest { activeUserId ->
-            activeUserId
-                ?.let { userId ->
-                    authDiskSource
-                        .getAccountTokensFlow(userId)
-                        .map { accountTokens ->
-                            accountTokens
-                                ?.accessToken
-                                ?.let { AuthState.Authenticated(it) }
-                                ?: AuthState.Unauthenticated
-                        }
-                }
-                ?: flowOf(AuthState.Unauthenticated)
-        }
-        .stateIn(
-            scope = unconfinedScope,
-            started = SharingStarted.Eagerly,
-            initialValue = AuthState.Uninitialized,
-        )
-
     private val duoTokenChannel = Channel<DuoCallbackTokenResult>(capacity = Int.MAX_VALUE)
     override val duoTokenResultFlow: Flow<DuoCallbackTokenResult> = duoTokenChannel.receiveAsFlow()
 
@@ -285,13 +254,6 @@ internal class AuthRepositoryImpl(
             }
         }
 
-    override val organizations: List<Organization>
-        get() = activeUserId
-            ?.let { authDiskSource.getOrganizations(it) }
-            ?.filter { it.status == OrganizationStatusType.CONFIRMED }
-            .orEmpty()
-            .toOrganizations()
-
     override val showWelcomeCarousel: Boolean
         get() = !settingsRepository.hasUserLoggedInOrCreatedAccount
 
@@ -1530,20 +1492,6 @@ internal class AuthRepositoryImpl(
         }
     }
 
-    override suspend fun leaveOrganization(organizationId: String): LeaveOrganizationResult =
-        organizationService.leaveOrganization(organizationId).fold(
-            onSuccess = { LeaveOrganizationResult.Success },
-            onFailure = { LeaveOrganizationResult.Error(error = it) },
-        )
-
-    override suspend fun revokeFromOrganization(
-        organizationId: String,
-    ): RevokeFromOrganizationResult =
-        organizationService.revokeFromOrganization(organizationId).fold(
-            onSuccess = { RevokeFromOrganizationResult.Success },
-            onFailure = { RevokeFromOrganizationResult.Error(error = it) },
-        )
-
     /**
      * Enrolls the active user in password reset if their organization requires it.
      */

app/src/main/kotlin/com/x8bit/bitwarden/data/auth/repository/di/AuthRepositoryModule.kt

@@ -11,26 +11,24 @@ import com.bitwarden.network.service.OrganizationService
 import com.x8bit.bitwarden.data.auth.datasource.disk.AuthDiskSource
 import com.x8bit.bitwarden.data.auth.datasource.sdk.AuthSdkSource
 import com.x8bit.bitwarden.data.auth.manager.AuthRequestManager
+import com.x8bit.bitwarden.data.auth.manager.AuthStateManager
 import com.x8bit.bitwarden.data.auth.manager.KdfManager
 import com.x8bit.bitwarden.data.auth.manager.KeyConnectorManager
+import com.x8bit.bitwarden.data.auth.manager.OrganizationManager
 import com.x8bit.bitwarden.data.auth.manager.TrustedDeviceManager
 import com.x8bit.bitwarden.data.auth.manager.UserLogoutManager
 import com.x8bit.bitwarden.data.auth.manager.UserStateManager
-import com.x8bit.bitwarden.data.auth.manager.UserStateManagerImpl
 import com.x8bit.bitwarden.data.auth.repository.AuthRepository
 import com.x8bit.bitwarden.data.auth.repository.AuthRepositoryImpl
 import com.x8bit.bitwarden.data.platform.datasource.disk.SettingsDiskSource
 import com.x8bit.bitwarden.data.platform.manager.BiometricsEncryptionManager
 import com.x8bit.bitwarden.data.platform.manager.FeatureFlagManager
-import com.x8bit.bitwarden.data.platform.manager.FirstTimeActionManager
 import com.x8bit.bitwarden.data.platform.manager.LogsManager
-import com.x8bit.bitwarden.data.platform.manager.PolicyManager
 import com.x8bit.bitwarden.data.platform.manager.PushManager
 import com.x8bit.bitwarden.data.platform.manager.policy.PasswordPolicyManager
 import com.x8bit.bitwarden.data.platform.repository.EnvironmentRepository
 import com.x8bit.bitwarden.data.platform.repository.SettingsRepository
 import com.x8bit.bitwarden.data.vault.datasource.sdk.VaultSdkSource
-import com.x8bit.bitwarden.data.vault.manager.VaultLockManager
 import com.x8bit.bitwarden.data.vault.repository.VaultRepository
 import dagger.Module
 import dagger.Provides
@@ -76,6 +74,8 @@ object AuthRepositoryModule {
         kdfManager: KdfManager,
         toastManager: ToastManager,
         featureFlagManager: FeatureFlagManager,
+        organizationManager: OrganizationManager,
+        authStateManager: AuthStateManager,
     ): AuthRepository = AuthRepositoryImpl(
         clock = clock,
         accountsService = accountsService,
@@ -104,21 +104,7 @@ object AuthRepositoryModule {
         kdfManager = kdfManager,
         toastManager = toastManager,
         featureFlagManager = featureFlagManager,
-    )
-
-    @Provides
-    @Singleton
-    fun providesUserStateManager(
-        authDiskSource: AuthDiskSource,
-        firstTimeActionManager: FirstTimeActionManager,
-        vaultLockManager: VaultLockManager,
-        policyManager: PolicyManager,
-        dispatcherManager: DispatcherManager,
-    ): UserStateManager = UserStateManagerImpl(
-        authDiskSource = authDiskSource,
-        firstTimeActionManager = firstTimeActionManager,
-        vaultLockManager = vaultLockManager,
-        policyManager = policyManager,
-        dispatcherManager = dispatcherManager,
+        organizationManager = organizationManager,
+        authStateManager = authStateManager,
     )
 }

app/src/main/kotlin/com/x8bit/bitwarden/data/platform/manager/di/PlatformManagerModule.kt

@@ -21,6 +21,8 @@ import com.bitwarden.network.service.PushService
 import com.x8bit.bitwarden.data.auth.datasource.disk.AuthDiskSource
 import com.x8bit.bitwarden.data.auth.datasource.sdk.AuthSdkSource
 import com.x8bit.bitwarden.data.auth.manager.AddTotpItemFromAuthenticatorManager
+import com.x8bit.bitwarden.data.auth.manager.AuthStateManager
+import com.x8bit.bitwarden.data.auth.manager.OrganizationManager
 import com.x8bit.bitwarden.data.auth.repository.AuthRepository
 import com.x8bit.bitwarden.data.autofill.accessibility.manager.AccessibilityEnabledManager
 import com.x8bit.bitwarden.data.autofill.manager.AutofillEnabledManager
@@ -128,14 +130,18 @@ object PlatformManagerModule {
     @Provides
     @Singleton
     fun provideOrganizationEventManager(
-        authRepository: AuthRepository,
+        authStateManager: AuthStateManager,
+        organizationManager: OrganizationManager,
         vaultRepository: VaultRepository,
+        authDiskSource: AuthDiskSource,
         clock: Clock,
         dispatcherManager: DispatcherManager,
         eventDiskSource: EventDiskSource,
         eventService: EventService,
     ): OrganizationEventManager = OrganizationEventManagerImpl(
-        authRepository = authRepository,
+        authStateManager = authStateManager,
+        authDiskSource = authDiskSource,
+        organizationManager = organizationManager,
         vaultRepository = vaultRepository,
         clock = clock,
         dispatcherManager = dispatcherManager,

app/src/main/kotlin/com/x8bit/bitwarden/data/platform/manager/event/OrganizationEventManagerImpl.kt

@@ -6,7 +6,9 @@ import androidx.lifecycle.ProcessLifecycleOwner
 import com.bitwarden.core.data.manager.dispatcher.DispatcherManager
 import com.bitwarden.network.model.OrganizationEventJson
 import com.bitwarden.network.service.EventService
-import com.x8bit.bitwarden.data.auth.repository.AuthRepository
+import com.x8bit.bitwarden.data.auth.datasource.disk.AuthDiskSource
+import com.x8bit.bitwarden.data.auth.manager.AuthStateManager
+import com.x8bit.bitwarden.data.auth.manager.OrganizationManager
 import com.x8bit.bitwarden.data.auth.repository.model.AuthState
 import com.x8bit.bitwarden.data.platform.datasource.disk.EventDiskSource
 import com.x8bit.bitwarden.data.platform.manager.model.OrganizationEvent
@@ -34,9 +36,11 @@ private const val UPLOAD_DELAY_INTERVAL_MS: Long = 300_000L
  * Default implementation of [OrganizationEventManager].
  */
 @Suppress("LongParameterList")
-class OrganizationEventManagerImpl(
+internal class OrganizationEventManagerImpl(
     private val clock: Clock,
-    private val authRepository: AuthRepository,
+    private val authStateManager: AuthStateManager,
+    private val authDiskSource: AuthDiskSource,
+    private val organizationManager: OrganizationManager,
     private val vaultRepository: VaultRepository,
     private val eventDiskSource: EventDiskSource,
     private val eventService: EventService,
@@ -46,6 +50,8 @@ class OrganizationEventManagerImpl(
     private val ioScope = CoroutineScope(dispatcherManager.io)
     private var job: Job = Job().apply { complete() }
 
+    private val activeUserId: String? get() = authDiskSource.userState?.activeUserId
+
     init {
         processLifecycleOwner.lifecycle.addObserver(
             object : DefaultLifecycleObserver {
@@ -57,9 +63,9 @@ class OrganizationEventManagerImpl(
     }
 
     override fun trackEvent(event: OrganizationEvent) {
-        val userId = authRepository.activeUserId ?: return
-        if (authRepository.authStateFlow.value !is AuthState.Authenticated) return
-        val organizations = authRepository.organizations.filter { it.shouldUseEvents }
+        val userId = activeUserId ?: return
+        if (authStateManager.authStateFlow.value !is AuthState.Authenticated) return
+        val organizations = organizationManager.organizations.filter { it.shouldUseEvents }
         if (organizations.none()) return
 
         ioScope.launch {
@@ -85,7 +91,7 @@ class OrganizationEventManagerImpl(
     }
 
     private suspend fun uploadEvents() {
-        val userId = authRepository.activeUserId ?: return
+        val userId = activeUserId ?: return
         val events = eventDiskSource
             .getOrganizationEvents(userId = userId)
             .takeUnless { it.isEmpty() }

app/src/test/kotlin/com/x8bit/bitwarden/data/auth/manager/AuthStateManagerTests.kt

@@ -0,0 +1,162 @@
+package com.x8bit.bitwarden.data.auth.manager
+
+import app.cash.turbine.test
+import com.bitwarden.core.data.manager.dispatcher.FakeDispatcherManager
+import com.bitwarden.data.datasource.disk.model.EnvironmentUrlDataJson
+import com.bitwarden.network.model.KdfJson
+import com.bitwarden.network.model.KdfTypeJson
+import com.bitwarden.network.model.MasterPasswordUnlockDataJson
+import com.bitwarden.network.model.UserDecryptionOptionsJson
+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.auth.repository.model.AuthState
+import kotlinx.coroutines.test.runTest
+import org.junit.jupiter.api.Assertions.assertEquals
+import org.junit.jupiter.api.Test
+import java.time.Instant
+
+class AuthStateManagerTests {
+
+    private val fakeAuthDiskSource = FakeAuthDiskSource()
+
+    private val authStateManager: AuthStateManager = AuthStateManagerImpl(
+        authDiskSource = fakeAuthDiskSource,
+        dispatcherManager = FakeDispatcherManager(),
+    )
+
+    @Test
+    fun `authStateFlow should react to user state changes and account token changes`() = runTest {
+        authStateManager.authStateFlow.test {
+            assertEquals(AuthState.Unauthenticated, awaitItem())
+
+            // Store the tokens, nothing happens yet since there is technically no active user yet
+            fakeAuthDiskSource.storeAccountTokens(
+                userId = USER_ID_1,
+                accountTokens = ACCOUNT_TOKENS_1,
+            )
+            expectNoEvents()
+            // Update the active user, we are now authenticated
+            fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
+            assertEquals(AuthState.Authenticated(accessToken = ACCESS_TOKEN), awaitItem())
+
+            // Adding a tokens for the non-active user does not update the state
+            fakeAuthDiskSource.storeAccountTokens(
+                userId = USER_ID_2,
+                accountTokens = ACCOUNT_TOKENS_2,
+            )
+            expectNoEvents()
+            // Adding a non-active user does not update the state
+            fakeAuthDiskSource.userState = MULTI_USER_STATE
+            expectNoEvents()
+
+            // Changing the active users tokens causes an update
+            val newAccessToken = "new_access_token"
+            fakeAuthDiskSource.storeAccountTokens(
+                userId = USER_ID_1,
+                accountTokens = ACCOUNT_TOKENS_1.copy(accessToken = newAccessToken),
+            )
+            assertEquals(AuthState.Authenticated(newAccessToken), awaitItem())
+
+            // Change the active user causes an update
+            fakeAuthDiskSource.userState = MULTI_USER_STATE.copy(activeUserId = USER_ID_2)
+            assertEquals(AuthState.Authenticated(accessToken = ACCESS_TOKEN_2), awaitItem())
+
+            // Clearing the tokens of the active state results in the Unauthenticated state
+            fakeAuthDiskSource.storeAccountTokens(userId = USER_ID_2, accountTokens = null)
+            assertEquals(AuthState.Unauthenticated, awaitItem())
+        }
+    }
+}
+
+private const val USER_ID_1 = "2a135b23-e1fb-42c9-bec3-573857bc8181"
+private const val USER_ID_2 = "b9d32ec0-6497-4582-9798-b350f53bfa02"
+private const val EMAIL = "test@bitwarden.com"
+private const val EMAIL_2 = "test2@bitwarden.com"
+private const val ACCESS_TOKEN = "accessToken"
+private const val ACCESS_TOKEN_2 = "accessToken2"
+private const val REFRESH_TOKEN = "refreshToken"
+
+private val ACCOUNT_TOKENS_1: AccountTokensJson = AccountTokensJson(
+    accessToken = ACCESS_TOKEN,
+    refreshToken = REFRESH_TOKEN,
+)
+private val ACCOUNT_TOKENS_2: AccountTokensJson = AccountTokensJson(
+    accessToken = ACCESS_TOKEN_2,
+    refreshToken = "refreshToken",
+)
+
+private val ACCOUNT_1: AccountJson = AccountJson(
+    profile = AccountJson.Profile(
+        userId = USER_ID_1,
+        email = EMAIL,
+        isEmailVerified = true,
+        name = "Bitwarden Tester",
+        hasPremiumPersonally = false,
+        hasPremiumFromOrganization = null,
+        stamp = null,
+        organizationId = null,
+        avatarColorHex = null,
+        forcePasswordResetReason = null,
+        kdfType = KdfTypeJson.ARGON2_ID,
+        kdfIterations = 600000,
+        kdfMemory = 16,
+        kdfParallelism = 4,
+        isTwoFactorEnabled = false,
+        creationDate = Instant.parse("2024-09-13T01:00:00.00Z"),
+        userDecryptionOptions = UserDecryptionOptionsJson(
+            hasMasterPassword = true,
+            trustedDeviceUserDecryptionOptions = null,
+            keyConnectorUserDecryptionOptions = null,
+            masterPasswordUnlock = MasterPasswordUnlockDataJson(
+                kdf = KdfJson(
+                    kdfType = KdfTypeJson.ARGON2_ID,
+                    iterations = 600000,
+                    memory = 16,
+                    parallelism = 4,
+                ),
+                masterKeyWrappedUserKey = "encryptedUserKey",
+                salt = EMAIL,
+            ),
+        ),
+    ),
+    settings = AccountJson.Settings(
+        environmentUrlData = EnvironmentUrlDataJson.DEFAULT_US,
+    ),
+)
+private val ACCOUNT_2 = AccountJson(
+    profile = AccountJson.Profile(
+        userId = USER_ID_2,
+        email = EMAIL_2,
+        isEmailVerified = true,
+        name = "Bitwarden Tester 2",
+        hasPremiumPersonally = false,
+        hasPremiumFromOrganization = null,
+        stamp = null,
+        organizationId = null,
+        avatarColorHex = null,
+        forcePasswordResetReason = null,
+        kdfType = KdfTypeJson.PBKDF2_SHA256,
+        kdfIterations = 400000,
+        kdfMemory = null,
+        kdfParallelism = null,
+        userDecryptionOptions = null,
+        isTwoFactorEnabled = true,
+        creationDate = Instant.parse("2024-09-13T01:00:00.00Z"),
+    ),
+    settings = AccountJson.Settings(
+        environmentUrlData = EnvironmentUrlDataJson.DEFAULT_EU,
+    ),
+)
+private val SINGLE_USER_STATE_1: UserStateJson = UserStateJson(
+    activeUserId = USER_ID_1,
+    accounts = mapOf(USER_ID_1 to ACCOUNT_1),
+)
+private val MULTI_USER_STATE = UserStateJson(
+    activeUserId = USER_ID_1,
+    accounts = mapOf(
+        USER_ID_1 to ACCOUNT_1,
+        USER_ID_2 to ACCOUNT_2,
+    ),
+)

app/src/test/kotlin/com/x8bit/bitwarden/data/auth/manager/OrganizationManagerTests.kt

@@ -0,0 +1,163 @@
+package com.x8bit.bitwarden.data.auth.manager
+
+import com.bitwarden.core.data.util.asFailure
+import com.bitwarden.core.data.util.asSuccess
+import com.bitwarden.data.datasource.disk.model.EnvironmentUrlDataJson
+import com.bitwarden.network.model.KdfJson
+import com.bitwarden.network.model.KdfTypeJson
+import com.bitwarden.network.model.MasterPasswordUnlockDataJson
+import com.bitwarden.network.model.UserDecryptionOptionsJson
+import com.bitwarden.network.model.createMockOrganizationNetwork
+import com.bitwarden.network.service.OrganizationService
+import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountJson
+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.auth.repository.model.LeaveOrganizationResult
+import com.x8bit.bitwarden.data.auth.repository.model.Organization
+import com.x8bit.bitwarden.data.auth.repository.model.RevokeFromOrganizationResult
+import com.x8bit.bitwarden.data.auth.repository.model.createMockOrganization
+import io.mockk.coEvery
+import io.mockk.coVerify
+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.Instant
+
+class OrganizationManagerTests {
+
+    private val fakeAuthDiskSource = FakeAuthDiskSource()
+    private val organizationService: OrganizationService = mockk()
+
+    private val organizationManager: OrganizationManager = OrganizationManagerImpl(
+        authDiskSource = fakeAuthDiskSource,
+        organizationService = organizationService,
+    )
+
+    @Test
+    fun `organizations should return an empty list when there is no active user`() = runTest {
+        assertEquals(emptyList<Organization>(), organizationManager.organizations)
+    }
+
+    @Test
+    fun `organizations should pull from the organizations in the AuthDiskSource`() = runTest {
+        fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
+        fakeAuthDiskSource.storeOrganizations(
+            userId = USER_ID_1,
+            organizations = listOf(createMockOrganizationNetwork(number = 0)),
+        )
+        assertEquals(listOf(createMockOrganization(number = 0)), organizationManager.organizations)
+    }
+
+    @Test
+    @Suppress("MaxLineLength")
+    fun `leaveOrganization should return success when organizationService leaveOrganization succeeds`() =
+        runTest {
+            coEvery {
+                organizationService.leaveOrganization(organizationId = any())
+            } returns Unit.asSuccess()
+
+            val continueResult = organizationManager.leaveOrganization(organizationId = "mockId-1")
+
+            coVerify(exactly = 1) {
+                organizationService.leaveOrganization(organizationId = any())
+            }
+            assertEquals(LeaveOrganizationResult.Success, continueResult)
+        }
+
+    @Test
+    fun `leaveOrganization should return error when organizationService leaveOrganization fails`() =
+        runTest {
+            val error = Throwable("Fail")
+            coEvery {
+                organizationService.leaveOrganization(organizationId = any())
+            } returns error.asFailure()
+
+            val continueResult = organizationManager.leaveOrganization(organizationId = "mockId-1")
+
+            coVerify(exactly = 1) {
+                organizationService.leaveOrganization(organizationId = any())
+            }
+            assertEquals(LeaveOrganizationResult.Error(error = error), continueResult)
+        }
+
+    @Test
+    @Suppress("MaxLineLength")
+    fun `revokeFromOrganization should return success when organizationService revokeFromOrganization succeeds`() =
+        runTest {
+            coEvery {
+                organizationService.revokeFromOrganization(organizationId = any())
+            } returns Unit.asSuccess()
+
+            val result = organizationManager.revokeFromOrganization(organizationId = "mockId-1")
+
+            coVerify(exactly = 1) {
+                organizationService.revokeFromOrganization(organizationId = any())
+            }
+            assertEquals(RevokeFromOrganizationResult.Success, result)
+        }
+
+    @Test
+    @Suppress("MaxLineLength")
+    fun `revokeFromOrganization should return error when organizationService revokeFromOrganization fails`() =
+        runTest {
+            val error = Throwable("Fail")
+            coEvery {
+                organizationService.revokeFromOrganization(organizationId = any())
+            } returns error.asFailure()
+
+            val result = organizationManager.revokeFromOrganization(organizationId = "mockId-1")
+
+            coVerify(exactly = 1) {
+                organizationService.revokeFromOrganization(organizationId = any())
+            }
+            assertEquals(RevokeFromOrganizationResult.Error(error = error), result)
+        }
+}
+
+private const val USER_ID_1 = "2a135b23-e1fb-42c9-bec3-573857bc8181"
+private const val EMAIL = "test@bitwarden.com"
+
+private val PROFILE_1 = AccountJson.Profile(
+    userId = USER_ID_1,
+    email = EMAIL,
+    isEmailVerified = true,
+    name = "Bitwarden Tester",
+    hasPremiumPersonally = false,
+    hasPremiumFromOrganization = null,
+    stamp = null,
+    organizationId = null,
+    avatarColorHex = null,
+    forcePasswordResetReason = null,
+    kdfType = KdfTypeJson.ARGON2_ID,
+    kdfIterations = 600000,
+    kdfMemory = 16,
+    kdfParallelism = 4,
+    isTwoFactorEnabled = false,
+    creationDate = Instant.parse("2024-09-13T01:00:00.00Z"),
+    userDecryptionOptions = UserDecryptionOptionsJson(
+        hasMasterPassword = true,
+        trustedDeviceUserDecryptionOptions = null,
+        keyConnectorUserDecryptionOptions = null,
+        masterPasswordUnlock = MasterPasswordUnlockDataJson(
+            kdf = KdfJson(
+                kdfType = KdfTypeJson.ARGON2_ID,
+                iterations = 600000,
+                memory = 16,
+                parallelism = 4,
+            ),
+            masterKeyWrappedUserKey = "encryptedUserKey",
+            salt = EMAIL,
+        ),
+    ),
+)
+private val ACCOUNT_1 = AccountJson(
+    profile = PROFILE_1,
+    settings = AccountJson.Settings(
+        environmentUrlData = EnvironmentUrlDataJson.DEFAULT_US,
+    ),
+)
+private val SINGLE_USER_STATE_1 = UserStateJson(
+    activeUserId = USER_ID_1,
+    accounts = mapOf(USER_ID_1 to ACCOUNT_1),
+)

app/src/test/kotlin/com/x8bit/bitwarden/data/auth/repository/AuthRepositoryTest.kt

@@ -66,7 +66,6 @@ import com.bitwarden.network.model.VerifyEmailTokenRequestJson
 import com.bitwarden.network.model.VerifyEmailTokenResponseJson
 import com.bitwarden.network.model.createMockAccountKeysJson
 import com.bitwarden.network.model.createMockAccountKeysJsonWithNullFields
-import com.bitwarden.network.model.createMockOrganizationNetwork
 import com.bitwarden.network.service.AccountsService
 import com.bitwarden.network.service.DevicesService
 import com.bitwarden.network.service.HaveIBeenPwnedService
@@ -83,8 +82,10 @@ import com.x8bit.bitwarden.data.auth.datasource.disk.util.FakeAuthDiskSource
 import com.x8bit.bitwarden.data.auth.datasource.sdk.AuthSdkSource
 import com.x8bit.bitwarden.data.auth.datasource.sdk.util.toKdfRequestModel
 import com.x8bit.bitwarden.data.auth.manager.AuthRequestManager
+import com.x8bit.bitwarden.data.auth.manager.AuthStateManager
 import com.x8bit.bitwarden.data.auth.manager.KdfManager
 import com.x8bit.bitwarden.data.auth.manager.KeyConnectorManager
+import com.x8bit.bitwarden.data.auth.manager.OrganizationManager
 import com.x8bit.bitwarden.data.auth.manager.TrustedDeviceManager
 import com.x8bit.bitwarden.data.auth.manager.UserLogoutManager
 import com.x8bit.bitwarden.data.auth.manager.UserStateManager
@@ -92,18 +93,15 @@ import com.x8bit.bitwarden.data.auth.manager.model.AuthRequest
 import com.x8bit.bitwarden.data.auth.manager.model.MigrateExistingUserToKeyConnectorResult
 import com.x8bit.bitwarden.data.auth.manager.model.MigrateNewUserToKeyConnectorResult
 import com.x8bit.bitwarden.data.auth.repository.AuthRepositoryTest.Companion.MASTER_PASSWORD_POLICY_OPTIONS
-import com.x8bit.bitwarden.data.auth.repository.model.AuthState
 import com.x8bit.bitwarden.data.auth.repository.model.BreachCountResult
 import com.x8bit.bitwarden.data.auth.repository.model.DeleteAccountResult
 import com.x8bit.bitwarden.data.auth.repository.model.DeviceInfo
 import com.x8bit.bitwarden.data.auth.repository.model.EmailTokenResult
 import com.x8bit.bitwarden.data.auth.repository.model.GetDevicesResult
 import com.x8bit.bitwarden.data.auth.repository.model.KnownDeviceResult
-import com.x8bit.bitwarden.data.auth.repository.model.LeaveOrganizationResult
 import com.x8bit.bitwarden.data.auth.repository.model.LoginResult
 import com.x8bit.bitwarden.data.auth.repository.model.LogoutReason
 import com.x8bit.bitwarden.data.auth.repository.model.NewSsoUserResult
-import com.x8bit.bitwarden.data.auth.repository.model.Organization
 import com.x8bit.bitwarden.data.auth.repository.model.PasswordHintResult
 import com.x8bit.bitwarden.data.auth.repository.model.PolicyInformation
 import com.x8bit.bitwarden.data.auth.repository.model.PrevalidateSsoResult
@@ -284,6 +282,8 @@ class AuthRepositoryTest {
         every { getFeatureFlag(FlagKey.V2EncryptionTde) } returns true
         every { getFeatureFlag(FlagKey.V2EncryptionPassword) } returns true
     }
+    private val authStateManager: AuthStateManager = mockk()
+    private val organizationManager: OrganizationManager = mockk()
 
     private val repository: AuthRepository = AuthRepositoryImpl(
         clock = FIXED_CLOCK,
@@ -313,6 +313,8 @@ class AuthRepositoryTest {
         kdfManager = kdfManager,
         toastManager = toastManager,
         featureFlagManager = featureFlagManager,
+        organizationManager = organizationManager,
+        authStateManager = authStateManager,
     )
 
     @BeforeEach
@@ -345,52 +347,6 @@ class AuthRepositoryTest {
         )
     }
 
-    @Test
-    fun `authStateFlow should react to user state changes and account token changes`() = runTest {
-        repository.authStateFlow.test {
-            assertEquals(AuthState.Unauthenticated, awaitItem())
-
-            // Store the tokens, nothing happens yet since there is technically no active user yet
-            fakeAuthDiskSource.storeAccountTokens(
-                userId = USER_ID_1,
-                accountTokens = ACCOUNT_TOKENS_1,
-            )
-            expectNoEvents()
-            // Update the active user, we are now authenticated
-            fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
-            assertEquals(AuthState.Authenticated(ACCESS_TOKEN), awaitItem())
-
-            // Adding a tokens for the non-active user does not update the state
-            fakeAuthDiskSource.storeAccountTokens(
-                userId = USER_ID_2,
-                accountTokens = ACCOUNT_TOKENS_2,
-            )
-            expectNoEvents()
-            // Adding a non-active user does not update the state
-            fakeAuthDiskSource.userState = MULTI_USER_STATE
-            expectNoEvents()
-
-            // Changing the active users tokens causes an update
-            val newAccessToken = "new_access_token"
-            fakeAuthDiskSource.storeAccountTokens(
-                userId = USER_ID_1,
-                accountTokens = ACCOUNT_TOKENS_1.copy(accessToken = newAccessToken),
-            )
-            assertEquals(AuthState.Authenticated(newAccessToken), awaitItem())
-
-            // Change the active user causes an update
-            fakeAuthDiskSource.userState = MULTI_USER_STATE.copy(activeUserId = USER_ID_2)
-            assertEquals(AuthState.Authenticated(ACCESS_TOKEN_2), awaitItem())
-
-            // Clearing the tokens of the active state results in the Unauthenticated state
-            fakeAuthDiskSource.storeAccountTokens(
-                userId = USER_ID_2,
-                accountTokens = null,
-            )
-            assertEquals(AuthState.Unauthenticated, awaitItem())
-        }
-    }
-
     @Test
     fun `rememberedEmailAddress should pull from and update AuthDiskSource`() {
         // AuthDiskSource and the repository start with the same value.
@@ -449,21 +405,6 @@ class AuthRepositoryTest {
         assertEquals(false, repository.shouldTrustDevice)
     }
 
-    @Test
-    fun `organizations should return an empty list when there is no active user`() = runTest {
-        assertEquals(emptyList<Organization>(), repository.organizations)
-    }
-
-    @Test
-    fun `organizations should pull from the organizations in the AuthDiskSource`() = runTest {
-        fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
-        fakeAuthDiskSource.storeOrganizations(
-            userId = USER_ID_1,
-            organizations = listOf(createMockOrganizationNetwork(number = 0)),
-        )
-        assertEquals(listOf(createMockOrganization(number = 0)), repository.organizations)
-    }
-
     @Test
     fun `delete account fails if not logged in`() = runTest {
         val masterPassword = "hello world"
@@ -1692,7 +1633,6 @@ class AuthRepositoryTest {
         } returns error.asFailure()
         val result = repository.login(email = EMAIL, password = PASSWORD)
         assertEquals(LoginResult.Error(error = error), result)
-        assertEquals(AuthState.Unauthenticated, repository.authStateFlow.value)
         coVerify { identityService.preLogin(email = EMAIL) }
     }
 
@@ -1717,7 +1657,6 @@ class AuthRepositoryTest {
             } returns error.asFailure()
             val result = repository.login(email = EMAIL, password = PASSWORD)
             assertEquals(LoginResult.Error(error = error), result)
-            assertEquals(AuthState.Unauthenticated, repository.authStateFlow.value)
             coVerify { identityService.preLogin(email = EMAIL) }
             coVerify {
                 identityService.getToken(
@@ -1753,7 +1692,6 @@ class AuthRepositoryTest {
             } returns RuntimeException().asFailure()
             val result = repository.login(email = EMAIL, password = PASSWORD)
             assertEquals(LoginResult.UnofficialServerError, result)
-            assertEquals(AuthState.Unauthenticated, repository.authStateFlow.value)
             coVerify { identityService.preLogin(email = EMAIL) }
         }
 
@@ -1777,7 +1715,6 @@ class AuthRepositoryTest {
             } returns SSLHandshakeException("error").asFailure()
             val result = repository.login(email = EMAIL, password = PASSWORD)
             assertEquals(LoginResult.CertificateError, result)
-            assertEquals(AuthState.Unauthenticated, repository.authStateFlow.value)
             coVerify { identityService.preLogin(email = EMAIL) }
         }
 
@@ -1789,7 +1726,6 @@ class AuthRepositoryTest {
             } returns SSLHandshakeException("error").asFailure()
             val result = repository.login(email = EMAIL, password = PASSWORD)
             assertEquals(LoginResult.CertificateError, result)
-            assertEquals(AuthState.Unauthenticated, repository.authStateFlow.value)
             coVerify { identityService.preLogin(email = EMAIL) }
         }
 
@@ -1818,7 +1754,6 @@ class AuthRepositoryTest {
 
         val result = repository.login(email = EMAIL, password = PASSWORD)
         assertEquals(LoginResult.Error(errorMessage = "mock_error_message", error = null), result)
-        assertEquals(AuthState.Unauthenticated, repository.authStateFlow.value)
         coVerify { identityService.preLogin(email = EMAIL) }
         coVerify {
             identityService.getToken(
@@ -1863,7 +1798,6 @@ class AuthRepositoryTest {
                 LoginResult.NewDeviceVerification(errorMessage = "new device verification required"),
                 result,
             )
-            assertEquals(AuthState.Unauthenticated, repository.authStateFlow.value)
         }
 
     @Test
@@ -1908,7 +1842,6 @@ class AuthRepositoryTest {
             } returns SINGLE_USER_STATE_1
             val result = repository.login(email = EMAIL, password = PASSWORD)
             assertEquals(LoginResult.Success, result)
-            assertEquals(AuthState.Authenticated(ACCESS_TOKEN), repository.authStateFlow.value)
             coVerify { identityService.preLogin(email = EMAIL) }
             fakeAuthDiskSource.assertAccountCryptographicState(
                 userId = USER_ID_1,
@@ -2228,7 +2161,6 @@ class AuthRepositoryTest {
             } returns SINGLE_USER_STATE_1
             val result = repository.login(email = EMAIL, password = PASSWORD)
             assertEquals(LoginResult.Success, result)
-            assertEquals(AuthState.Authenticated(ACCESS_TOKEN), repository.authStateFlow.value)
             coVerify { identityService.preLogin(email = EMAIL) }
             fakeAuthDiskSource.assertAccountCryptographicState(
                 userId = USER_ID_1,
@@ -2322,7 +2254,6 @@ class AuthRepositoryTest {
                 LoginResult.Error(errorMessage = expectedErrorMessage, error = error),
                 result,
             )
-            assertEquals(AuthState.Unauthenticated, repository.authStateFlow.value)
             coVerify { identityService.preLogin(email = EMAIL) }
             fakeAuthDiskSource.assertAccountCryptographicState(
                 userId = USER_ID_1,
@@ -2409,7 +2340,6 @@ class AuthRepositoryTest {
                 password = PASSWORD,
             )
             assertEquals(LoginResult.Success, result)
-            assertEquals(AuthState.Authenticated(ACCESS_TOKEN), repository.authStateFlow.value)
             coVerify { identityService.preLogin(email = EMAIL) }
             fakeAuthDiskSource.assertMasterPasswordHash(
                 userId = USER_ID_1,
@@ -2509,7 +2439,6 @@ class AuthRepositoryTest {
             val result = repository.login(email = EMAIL, password = PASSWORD)
 
             assertEquals(LoginResult.Success, result)
-            assertEquals(AuthState.Authenticated(ACCESS_TOKEN), repository.authStateFlow.value)
             coVerify { identityService.preLogin(email = EMAIL) }
             fakeAuthDiskSource.assertAccountCryptographicState(
                 userId = USER_ID_1,
@@ -2581,7 +2510,6 @@ class AuthRepositoryTest {
                 ssoToken = null,
             ),
         )
-        assertEquals(AuthState.Unauthenticated, repository.authStateFlow.value)
         coVerify { identityService.preLogin(email = EMAIL) }
         coVerify {
             identityService.getToken(
@@ -2833,7 +2761,6 @@ class AuthRepositoryTest {
         } returns SINGLE_USER_STATE_1
         val result = repository.login(email = EMAIL, password = PASSWORD)
         assertEquals(LoginResult.Success, result)
-        assertEquals(AuthState.Authenticated(ACCESS_TOKEN), repository.authStateFlow.value)
         coVerify { identityService.preLogin(email = EMAIL) }
         fakeAuthDiskSource.assertAccountCryptographicState(
             userId = USER_ID_1,
@@ -2914,7 +2841,6 @@ class AuthRepositoryTest {
 
             val result = repository.login(email = EMAIL, password = PASSWORD)
             assertEquals(LoginResult.EncryptionKeyMigrationRequired, result)
-            assertEquals(AuthState.Unauthenticated, repository.authStateFlow.value)
             coVerify {
                 identityService.preLogin(email = EMAIL)
                 identityService.getToken(
@@ -2953,7 +2879,6 @@ class AuthRepositoryTest {
             masterPasswordHash = PASSWORD_HASH,
         )
         assertEquals(LoginResult.Error(error = error), result)
-        assertEquals(AuthState.Unauthenticated, repository.authStateFlow.value)
         coVerify {
             identityService.getToken(
                 email = EMAIL,
@@ -3002,7 +2927,6 @@ class AuthRepositoryTest {
                 LoginResult.Error(errorMessage = "mock_error_message", error = null),
                 result,
             )
-            assertEquals(AuthState.Unauthenticated, repository.authStateFlow.value)
             coVerify {
                 identityService.getToken(
                     email = EMAIL,
@@ -3067,7 +2991,6 @@ class AuthRepositoryTest {
                 masterPasswordHash = PASSWORD_HASH,
             )
             assertEquals(LoginResult.Success, result)
-            assertEquals(AuthState.Authenticated(ACCESS_TOKEN), repository.authStateFlow.value)
             fakeAuthDiskSource.assertAccountCryptographicState(
                 userId = USER_ID_1,
                 accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE_V2,
@@ -3160,7 +3083,6 @@ class AuthRepositoryTest {
                 masterPasswordHash = PASSWORD_HASH,
             )
             assertEquals(LoginResult.Success, result)
-            assertEquals(AuthState.Authenticated(ACCESS_TOKEN), repository.authStateFlow.value)
             fakeAuthDiskSource.assertAccountCryptographicState(
                 userId = USER_ID_1,
                 accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE_V2,
@@ -3240,7 +3162,6 @@ class AuthRepositoryTest {
                     ssoToken = null,
                 ),
             )
-            assertEquals(AuthState.Unauthenticated, repository.authStateFlow.value)
             coVerify {
                 identityService.getToken(
                     email = EMAIL,
@@ -3380,7 +3301,6 @@ class AuthRepositoryTest {
             organizationIdentifier = ORGANIZATION_IDENTIFIER,
         )
         assertEquals(LoginResult.Error(error = error), result)
-        assertEquals(AuthState.Unauthenticated, repository.authStateFlow.value)
         coVerify {
             identityService.getToken(
                 email = EMAIL,
@@ -3424,7 +3344,6 @@ class AuthRepositoryTest {
             organizationIdentifier = ORGANIZATION_IDENTIFIER,
         )
         assertEquals(LoginResult.Error(errorMessage = "mock_error_message", error = null), result)
-        assertEquals(AuthState.Unauthenticated, repository.authStateFlow.value)
         coVerify {
             identityService.getToken(
                 email = EMAIL,
@@ -3472,7 +3391,6 @@ class AuthRepositoryTest {
                 organizationIdentifier = ORGANIZATION_IDENTIFIER,
             )
             assertEquals(LoginResult.Success, result)
-            assertEquals(AuthState.Authenticated(ACCESS_TOKEN), repository.authStateFlow.value)
             fakeAuthDiskSource.assertAccountCryptographicState(
                 userId = USER_ID_1,
                 accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE_V2,
@@ -3541,7 +3459,6 @@ class AuthRepositoryTest {
                 organizationIdentifier = ORGANIZATION_IDENTIFIER,
             )
             assertEquals(LoginResult.Success, result)
-            assertEquals(AuthState.Authenticated(ACCESS_TOKEN), repository.authStateFlow.value)
             fakeAuthDiskSource.assertAccountCryptographicState(
                 userId = USER_ID_1,
                 accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE_V2,
@@ -3706,7 +3623,6 @@ class AuthRepositoryTest {
             )
 
             assertEquals(LoginResult.Success, result)
-            assertEquals(AuthState.Authenticated(ACCESS_TOKEN), repository.authStateFlow.value)
             fakeAuthDiskSource.assertAccountCryptographicState(
                 userId = USER_ID_1,
                 accountCryptographicState = accountCryptographicState,
@@ -3800,7 +3716,6 @@ class AuthRepositoryTest {
             )
 
             assertEquals(LoginResult.Success, result)
-            assertEquals(AuthState.Authenticated(ACCESS_TOKEN), repository.authStateFlow.value)
             fakeAuthDiskSource.assertAccountCryptographicState(
                 userId = USER_ID_1,
                 accountCryptographicState = accountCryptographicState,
@@ -4013,7 +3928,6 @@ class AuthRepositoryTest {
                 email = EMAIL,
             )
             assertEquals(LoginResult.Success, continueResult)
-            assertEquals(AuthState.Authenticated(ACCESS_TOKEN), repository.authStateFlow.value)
             fakeAuthDiskSource.assertAccountCryptographicState(
                 userId = USER_ID_1,
                 accountCryptographicState = keyConnectorResult.accountCryptographicState,
@@ -4192,7 +4106,6 @@ class AuthRepositoryTest {
                 email = EMAIL,
             )
             assertEquals(LoginResult.Success, result)
-            assertEquals(AuthState.Authenticated(ACCESS_TOKEN), repository.authStateFlow.value)
             fakeAuthDiskSource.assertAccountCryptographicState(
                 userId = USER_ID_1,
                 accountCryptographicState = keyConnectorResult.accountCryptographicState,
@@ -4266,7 +4179,6 @@ class AuthRepositoryTest {
                 masterPasswordHash = null,
             )
             assertEquals(LoginResult.Success, result)
-            assertEquals(AuthState.Authenticated(ACCESS_TOKEN), repository.authStateFlow.value)
             fakeAuthDiskSource.assertAccountCryptographicState(
                 userId = USER_ID_1,
                 accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE_V2,
@@ -4347,7 +4259,6 @@ class AuthRepositoryTest {
             )
 
             assertEquals(LoginResult.Success, result)
-            assertEquals(AuthState.Authenticated(ACCESS_TOKEN), repository.authStateFlow.value)
             fakeAuthDiskSource.assertAccountCryptographicState(
                 userId = USER_ID_1,
                 accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE_V2,
@@ -4434,8 +4345,6 @@ class AuthRepositoryTest {
             )
 
             assertEquals(LoginResult.Success, result)
-            assertEquals(AuthState.Authenticated(ACCESS_TOKEN), repository.authStateFlow.value)
-
             fakeAuthDiskSource.assertAccountCryptographicState(
                 userId = USER_ID_1,
                 accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE_V2,
@@ -4538,7 +4447,6 @@ class AuthRepositoryTest {
             )
 
             assertEquals(LoginResult.Success, result)
-            assertEquals(AuthState.Authenticated(ACCESS_TOKEN), repository.authStateFlow.value)
             fakeAuthDiskSource.assertAccountCryptographicState(
                 userId = USER_ID_1,
                 accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE_V2,
@@ -4614,7 +4522,6 @@ class AuthRepositoryTest {
             )
 
             assertEquals(LoginResult.Success, result)
-            assertEquals(AuthState.Authenticated(ACCESS_TOKEN), repository.authStateFlow.value)
             fakeAuthDiskSource.storeAccountCryptographicState(
                 userId = USER_ID_1,
                 accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE_V2,
@@ -4681,7 +4588,6 @@ class AuthRepositoryTest {
                     twoFactorProviders = null,
                 ),
             )
-            assertEquals(AuthState.Unauthenticated, repository.authStateFlow.value)
             coVerify {
                 identityService.getToken(
                     email = EMAIL,
@@ -4820,7 +4726,6 @@ class AuthRepositoryTest {
             organizationIdentifier = ORGANIZATION_IDENTIFIER,
         )
         assertEquals(LoginResult.Success, result)
-        assertEquals(AuthState.Authenticated(ACCESS_TOKEN), repository.authStateFlow.value)
         fakeAuthDiskSource.storeAccountCryptographicState(
             userId = USER_ID_1,
             accountCryptographicState = ACCOUNT_CRYPTOGRAPHIC_STATE_V2,
@@ -5000,14 +4905,14 @@ class AuthRepositoryTest {
     fun `removePassword with no keyConnectorUrl should return error`() = runTest {
         fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
         val organizations = listOf(
-            createMockOrganizationNetwork(
+            createMockOrganization(
                 number = 1,
-                isKeyConnectorEnabled = true,
-                type = OrganizationType.USER,
+                shouldUseKeyConnector = true,
+                role = OrganizationType.USER,
                 keyConnectorUrl = null,
             ),
         )
-        fakeAuthDiskSource.storeOrganizations(userId = USER_ID_1, organizations = organizations)
+        every { organizationManager.organizations } returns organizations
 
         val result = repository.removePassword(masterPassword = PASSWORD)
 
@@ -5024,14 +4929,14 @@ class AuthRepositoryTest {
             val url = "www.example.com"
             val error = Throwable("Fail!")
             val organizations = listOf(
-                createMockOrganizationNetwork(
+                createMockOrganization(
                     number = 1,
-                    isKeyConnectorEnabled = true,
-                    type = OrganizationType.USER,
+                    shouldUseKeyConnector = true,
+                    role = OrganizationType.USER,
                     keyConnectorUrl = url,
                 ),
             )
-            fakeAuthDiskSource.storeOrganizations(userId = USER_ID_1, organizations = organizations)
+            every { organizationManager.organizations } returns organizations
             coEvery {
                 keyConnectorManager.migrateExistingUserToKeyConnector(
                     userId = USER_ID_1,
@@ -5056,14 +4961,14 @@ class AuthRepositoryTest {
             val error = Throwable("Fail!")
             val expectedResult = MigrateExistingUserToKeyConnectorResult.Error(error)
             val organizations = listOf(
-                createMockOrganizationNetwork(
+                createMockOrganization(
                     number = 1,
-                    isKeyConnectorEnabled = true,
-                    type = OrganizationType.USER,
+                    shouldUseKeyConnector = true,
+                    role = OrganizationType.USER,
                     keyConnectorUrl = url,
                 ),
             )
-            fakeAuthDiskSource.storeOrganizations(userId = USER_ID_1, organizations = organizations)
+            every { organizationManager.organizations } returns organizations
             coEvery {
                 keyConnectorManager.migrateExistingUserToKeyConnector(
                     userId = USER_ID_1,
@@ -5091,14 +4996,14 @@ class AuthRepositoryTest {
             val url = "www.example.com"
             val expectedResult = MigrateExistingUserToKeyConnectorResult.WrongPasswordError
             val organizations = listOf(
-                createMockOrganizationNetwork(
+                createMockOrganization(
                     number = 1,
-                    isKeyConnectorEnabled = true,
-                    type = OrganizationType.USER,
+                    shouldUseKeyConnector = true,
+                    role = OrganizationType.USER,
                     keyConnectorUrl = url,
                 ),
             )
-            fakeAuthDiskSource.storeOrganizations(userId = USER_ID_1, organizations = organizations)
+            every { organizationManager.organizations } returns organizations
             coEvery {
                 keyConnectorManager.migrateExistingUserToKeyConnector(
                     userId = USER_ID_1,
@@ -5125,14 +5030,14 @@ class AuthRepositoryTest {
             fakeAuthDiskSource.userState = SINGLE_USER_STATE_1
             val url = "www.example.com"
             val organizations = listOf(
-                createMockOrganizationNetwork(
+                createMockOrganization(
                     number = 1,
-                    isKeyConnectorEnabled = true,
-                    type = OrganizationType.USER,
+                    shouldUseKeyConnector = true,
+                    role = OrganizationType.USER,
                     keyConnectorUrl = url,
                 ),
             )
-            fakeAuthDiskSource.storeOrganizations(userId = USER_ID_1, organizations = organizations)
+            every { organizationManager.organizations } returns organizations
             coEvery {
                 keyConnectorManager.migrateExistingUserToKeyConnector(
                     userId = USER_ID_1,
@@ -7185,7 +7090,6 @@ class AuthRepositoryTest {
             every { settingsRepository.getUserHasLoggedInValue(USER_ID_1) } returns false
             val result = repository.login(email = EMAIL, password = PASSWORD)
             assertEquals(LoginResult.Success, result)
-            assertEquals(AuthState.Authenticated(ACCESS_TOKEN), repository.authStateFlow.value)
             coVerify { identityService.preLogin(email = EMAIL) }
             fakeAuthDiskSource.assertAccountCryptographicState(
                 userId = USER_ID_1,
@@ -7244,7 +7148,6 @@ class AuthRepositoryTest {
             every { settingsRepository.getUserHasLoggedInValue(USER_ID_1) } returns true
             val result = repository.login(email = EMAIL, password = PASSWORD)
             assertEquals(LoginResult.Success, result)
-            assertEquals(AuthState.Authenticated(ACCESS_TOKEN), repository.authStateFlow.value)
             coVerify { identityService.preLogin(email = EMAIL) }
             fakeAuthDiskSource.assertAccountCryptographicState(
                 userId = USER_ID_1,
@@ -7280,41 +7183,6 @@ class AuthRepositoryTest {
             )
         }
 
-    @Test
-    @Suppress("MaxLineLength")
-    fun `leaveOrganization should return success when organizationService leaveOrganization succeeds`() =
-        runTest {
-            coEvery {
-                organizationService.leaveOrganization(any())
-            } returns Unit.asSuccess()
-
-            val continueResult = repository.leaveOrganization("mockId-1")
-            coVerify {
-                organizationService.leaveOrganization(any())
-            }
-            assertEquals(
-                LeaveOrganizationResult.Success, continueResult,
-            )
-        }
-
-    @Test
-    fun `leaveOrganization should return error when organizationService leaveOrganization fails`() =
-        runTest {
-            val error = Throwable("Fail")
-            coEvery {
-                organizationService.leaveOrganization(any())
-            } returns error.asFailure()
-
-            val continueResult = repository.leaveOrganization("mockId-1")
-            coVerify {
-                organizationService.leaveOrganization(any())
-            }
-            assertEquals(
-                LeaveOrganizationResult.Error(error = error),
-                continueResult,
-            )
-        }
-
     companion object {
         private val FIXED_CLOCK: Clock = Clock.fixed(
             Instant.parse("2023-10-27T12:00:00Z"),

app/src/test/kotlin/com/x8bit/bitwarden/data/platform/manager/event/OrganizationEventManagerTest.kt

@@ -9,7 +9,10 @@ import com.bitwarden.network.model.OrganizationEventJson
 import com.bitwarden.network.model.OrganizationEventType
 import com.bitwarden.network.service.EventService
 import com.bitwarden.vault.CipherView
-import com.x8bit.bitwarden.data.auth.repository.AuthRepository
+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.auth.manager.AuthStateManager
+import com.x8bit.bitwarden.data.auth.manager.OrganizationManager
 import com.x8bit.bitwarden.data.auth.repository.model.AuthState
 import com.x8bit.bitwarden.data.auth.repository.model.createMockOrganization
 import com.x8bit.bitwarden.data.platform.datasource.disk.EventDiskSource
@@ -40,11 +43,6 @@ class OrganizationEventManagerTest {
     private val dispatcher = StandardTestDispatcher()
     private val fakeDispatcherManager = FakeDispatcherManager(io = dispatcher)
     private val mutableAuthStateFlow = MutableStateFlow<AuthState>(value = AuthState.Uninitialized)
-    private val authRepository = mockk<AuthRepository> {
-        every { activeUserId } returns USER_ID
-        every { authStateFlow } returns mutableAuthStateFlow
-        every { organizations } returns emptyList()
-    }
     private val mutableVaultItemStateFlow = MutableStateFlow<DataState<CipherView?>>(
         value = DataState.Loading,
     )
@@ -55,12 +53,26 @@ class OrganizationEventManagerTest {
     private val eventDiskSource = mockk<EventDiskSource> {
         coEvery { addOrganizationEvent(userId = any(), event = any()) } just runs
     }
+    private val authStateManager: AuthStateManager = mockk {
+        every { authStateFlow } returns mutableAuthStateFlow
+    }
+    private val fakeAuthDiskSource = FakeAuthDiskSource().apply {
+        userState = UserStateJson(
+            activeUserId = USER_ID,
+            accounts = mapOf(USER_ID to mockk()),
+        )
+    }
+    private val organizationManager: OrganizationManager = mockk {
+        every { organizations } returns emptyList()
+    }
 
     private val organizationEventManager: OrganizationEventManager = OrganizationEventManagerImpl(
         processLifecycleOwner = fakeLifecycleOwner,
         clock = fixedClock,
         dispatcherManager = fakeDispatcherManager,
-        authRepository = authRepository,
+        authDiskSource = fakeAuthDiskSource,
+        authStateManager = authStateManager,
+        organizationManager = organizationManager,
         vaultRepository = vaultRepository,
         eventService = eventService,
         eventDiskSource = eventDiskSource,
@@ -124,7 +136,7 @@ class OrganizationEventManagerTest {
 
     @Test
     fun `trackEvent should do nothing if there is no active user`() {
-        every { authRepository.activeUserId } returns null
+        fakeAuthDiskSource.userState = null
 
         organizationEventManager.trackEvent(
             event = OrganizationEvent.CipherClientAutoFilled(
@@ -154,7 +166,7 @@ class OrganizationEventManagerTest {
     fun `trackEvent should do nothing if the active user has no organizations that use events`() {
         mutableAuthStateFlow.value = AuthState.Authenticated(accessToken = "access-token")
         val organization = createMockOrganization(number = 1)
-        every { authRepository.organizations } returns listOf(organization)
+        every { organizationManager.organizations } returns listOf(organization)
 
         organizationEventManager.trackEvent(
             event = OrganizationEvent.CipherClientAutoFilled(
@@ -172,7 +184,7 @@ class OrganizationEventManagerTest {
     fun `trackEvent should do nothing if the cipher does not belong to an organization that uses events`() {
         mutableAuthStateFlow.value = AuthState.Authenticated(accessToken = "access-token")
         val organization = createMockOrganization(number = 1, shouldUseEvents = true)
-        every { authRepository.organizations } returns listOf(organization)
+        every { organizationManager.organizations } returns listOf(organization)
         val cipherView = createMockCipherView(number = 1)
         mutableVaultItemStateFlow.value = DataState.Loaded(data = cipherView)
 
@@ -195,7 +207,7 @@ class OrganizationEventManagerTest {
             id = "mockOrganizationId-1",
             shouldUseEvents = true,
         )
-        every { authRepository.organizations } returns listOf(organization)
+        every { organizationManager.organizations } returns listOf(organization)
         val cipherView = createMockCipherView(number = 1)
         mutableVaultItemStateFlow.value = DataState.Loaded(data = cipherView)