Save the current environment to UserStateJson when logging in (#260)

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

@@ -5,6 +5,7 @@ import com.bitwarden.core.Kdf
 import com.bitwarden.core.RegisterKeyResponse
 import com.bitwarden.core.RsaKeyPair
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountJson
+import com.x8bit.bitwarden.data.auth.datasource.disk.model.EnvironmentUrlDataJson
 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.datasource.network.model.GetTokenResponseJson
@@ -31,6 +32,8 @@ import com.x8bit.bitwarden.data.auth.repository.util.toUserState
 import com.x8bit.bitwarden.data.auth.util.toSdkParams
 import com.x8bit.bitwarden.data.platform.base.FakeDispatcherManager
 import com.x8bit.bitwarden.data.platform.manager.dispatcher.DispatcherManager
+import com.x8bit.bitwarden.data.platform.repository.model.Environment
+import com.x8bit.bitwarden.data.platform.repository.util.FakeEnvironmentRepository
 import com.x8bit.bitwarden.data.platform.util.asFailure
 import com.x8bit.bitwarden.data.platform.util.asSuccess
 import com.x8bit.bitwarden.data.vault.repository.VaultRepository
@@ -59,6 +62,11 @@ class AuthRepositoryTest {
     private val haveIBeenPwnedService: HaveIBeenPwnedService = mockk()
     private val vaultRepository: VaultRepository = mockk()
     private val fakeAuthDiskSource = FakeAuthDiskSource()
+    private val fakeEnvironmentRepository =
+        FakeEnvironmentRepository()
+            .apply {
+                environment = Environment.Us
+            }
     private val authSdkSource = mockk<AuthSdkSource> {
         coEvery {
             hashPassword(
@@ -91,8 +99,9 @@ class AuthRepositoryTest {
         haveIBeenPwnedService = haveIBeenPwnedService,
         authSdkSource = authSdkSource,
         authDiskSource = fakeAuthDiskSource,
-        dispatcherManager = dispatcherManager,
+        environmentRepository = fakeEnvironmentRepository,
         vaultRepository = vaultRepository,
+        dispatcherManager = dispatcherManager,
     )
 
     @BeforeEach
@@ -278,7 +287,10 @@ class AuthRepositoryTest {
                 vaultRepository.unlockVaultAndSync(masterPassword = PASSWORD)
             } returns VaultUnlockResult.Success
             every {
-                GET_TOKEN_RESPONSE_SUCCESS.toUserState(previousUserState = null)
+                GET_TOKEN_RESPONSE_SUCCESS.toUserState(
+                    previousUserState = null,
+                    environmentUrlData = EnvironmentUrlDataJson.DEFAULT_US,
+                )
             } returns SINGLE_USER_STATE_1
             val result = repository.login(email = EMAIL, password = PASSWORD, captchaToken = null)
             assertEquals(LoginResult.Success, result)
@@ -684,7 +696,10 @@ class AuthRepositoryTest {
             vaultRepository.unlockVaultAndSync(masterPassword = PASSWORD)
         } returns VaultUnlockResult.Success
         every {
-            GET_TOKEN_RESPONSE_SUCCESS.toUserState(previousUserState = null)
+            GET_TOKEN_RESPONSE_SUCCESS.toUserState(
+                previousUserState = null,
+                environmentUrlData = EnvironmentUrlDataJson.DEFAULT_US,
+            )
         } returns SINGLE_USER_STATE_1
 
         repository.login(email = EMAIL, password = PASSWORD, captchaToken = null)
@@ -733,7 +748,10 @@ class AuthRepositoryTest {
                 vaultRepository.unlockVaultAndSync(masterPassword = PASSWORD)
             } returns VaultUnlockResult.Success
             every {
-                GET_TOKEN_RESPONSE_SUCCESS.toUserState(previousUserState = SINGLE_USER_STATE_2)
+                GET_TOKEN_RESPONSE_SUCCESS.toUserState(
+                    previousUserState = SINGLE_USER_STATE_2,
+                    environmentUrlData = EnvironmentUrlDataJson.DEFAULT_US,
+                )
             } returns MULTI_USER_STATE
 
             repository.login(email = EMAIL, password = PASSWORD, captchaToken = null)

app/src/test/java/com/x8bit/bitwarden/data/auth/repository/util/GetTokenResponseExtensionsTest.kt

@@ -1,6 +1,7 @@
 package com.x8bit.bitwarden.data.auth.repository.util
 
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountJson
+import com.x8bit.bitwarden.data.auth.datasource.disk.model.EnvironmentUrlDataJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.UserStateJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.GetTokenResponseJson
 import com.x8bit.bitwarden.data.auth.datasource.network.model.KdfTypeJson
@@ -31,7 +32,10 @@ class GetTokenResponseExtensionsTest {
 
         assertEquals(
             SINGLE_USER_STATE_1,
-            GET_TOKEN_RESPONSE_SUCCESS.toUserState(previousUserState = null),
+            GET_TOKEN_RESPONSE_SUCCESS.toUserState(
+                previousUserState = null,
+                environmentUrlData = EnvironmentUrlDataJson.DEFAULT_US,
+            ),
         )
     }
 
@@ -41,7 +45,10 @@ class GetTokenResponseExtensionsTest {
 
         assertEquals(
             MULTI_USER_STATE,
-            GET_TOKEN_RESPONSE_SUCCESS.toUserState(previousUserState = SINGLE_USER_STATE_2),
+            GET_TOKEN_RESPONSE_SUCCESS.toUserState(
+                previousUserState = SINGLE_USER_STATE_2,
+                environmentUrlData = EnvironmentUrlDataJson.DEFAULT_US,
+            ),
         )
     }
 }
@@ -102,7 +109,7 @@ private val ACCOUNT_1 = AccountJson(
         refreshToken = "refreshToken",
     ),
     settings = AccountJson.Settings(
-        environmentUrlData = null,
+        environmentUrlData = EnvironmentUrlDataJson.DEFAULT_US,
     ),
 )
 private val ACCOUNT_2 = AccountJson(
@@ -127,7 +134,7 @@ private val ACCOUNT_2 = AccountJson(
         refreshToken = "refreshToken",
     ),
     settings = AccountJson.Settings(
-        environmentUrlData = null,
+        environmentUrlData = EnvironmentUrlDataJson.DEFAULT_US,
     ),
 )
 private val SINGLE_USER_STATE_1 = UserStateJson(

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

@@ -1,7 +1,10 @@
 package com.x8bit.bitwarden.data.platform.repository
 
 import app.cash.turbine.test
+import com.x8bit.bitwarden.data.auth.datasource.disk.model.AccountJson
 import com.x8bit.bitwarden.data.auth.datasource.disk.model.EnvironmentUrlDataJson
+import com.x8bit.bitwarden.data.auth.datasource.disk.model.UserStateJson
+import com.x8bit.bitwarden.data.auth.datasource.disk.util.FakeAuthDiskSource
 import com.x8bit.bitwarden.data.platform.base.FakeDispatcherManager
 import com.x8bit.bitwarden.data.platform.datasource.disk.EnvironmentDiskSource
 import com.x8bit.bitwarden.data.platform.manager.dispatcher.DispatcherManager
@@ -26,9 +29,11 @@ class EnvironmentRepositoryTest {
     private val dispatcherManager: DispatcherManager = FakeDispatcherManager()
 
     private val fakeEnvironmentDiskSource = FakeEnvironmentDiskSource()
+    private val fakeAuthDiskSource = FakeAuthDiskSource()
 
     private val repository = EnvironmentRepositoryImpl(
         environmentDiskSource = fakeEnvironmentDiskSource,
+        authDiskSource = fakeAuthDiskSource,
         dispatcherManager = dispatcherManager,
     )
 
@@ -42,6 +47,46 @@ class EnvironmentRepositoryTest {
         unmockkStatic(ENVIRONMENT_EXTENSIONS_PATH)
     }
 
+    @Test
+    fun `changes to the active user should update the environment if necessary`() {
+        assertEquals(
+            Environment.Us,
+            repository.environment,
+        )
+        assertEquals(
+            null,
+            fakeEnvironmentDiskSource.preAuthEnvironmentUrlData,
+        )
+
+        // Updating the environment for the active user to a non-null value triggers an update
+        // in the saved environment.
+        fakeAuthDiskSource.userState = getMockUserState(
+            environmentForActiveUser = EnvironmentUrlDataJson.DEFAULT_EU,
+        )
+        assertEquals(
+            Environment.Eu,
+            repository.environment,
+        )
+        assertEquals(
+            EnvironmentUrlDataJson.DEFAULT_EU,
+            fakeEnvironmentDiskSource.preAuthEnvironmentUrlData,
+        )
+
+        // Updating the environment for the active user to a null value leaves the current
+        // environment unchanged.
+        fakeAuthDiskSource.userState = getMockUserState(
+            environmentForActiveUser = null,
+        )
+        assertEquals(
+            Environment.Eu,
+            repository.environment,
+        )
+        assertEquals(
+            EnvironmentUrlDataJson.DEFAULT_EU,
+            fakeEnvironmentDiskSource.preAuthEnvironmentUrlData,
+        )
+    }
+
     @Test
     fun `environment should pull from and update EnvironmentDiskSource`() {
         val environmentUrlDataJson = mockk<EnvironmentUrlDataJson>()
@@ -101,6 +146,22 @@ class EnvironmentRepositoryTest {
             assertEquals(environment, awaitItem())
         }
     }
+
+    private fun getMockUserState(
+        environmentForActiveUser: EnvironmentUrlDataJson?,
+    ): UserStateJson =
+        UserStateJson(
+            activeUserId = "activeUserId",
+            accounts = mapOf(
+                "activeUserId" to AccountJson(
+                    profile = mockk(),
+                    tokens = mockk(),
+                    settings = AccountJson.Settings(
+                        environmentUrlData = environmentForActiveUser,
+                    ),
+                ),
+            ),
+        )
 }
 
 private const val ENVIRONMENT_EXTENSIONS_PATH =