feat: update teamId after migrating from personal to team account (WPB-12187)

data/src/commonMain/kotlin/com/wire/kalium/logic/data/user/CreateUserTeam.kt

@@ -1,3 +1,3 @@
 package com.wire.kalium.logic.data.user
 
-data class CreateUserTeam(val teamName: String)
+data class CreateUserTeam(val teamId: String, val teamName: String)

logic/src/commonMain/kotlin/com/wire/kalium/logic/data/user/UserRepository.kt

@@ -166,6 +166,7 @@ interface UserRepository {
     suspend fun getUsersMinimizedByQualifiedIDs(userIds: List<UserId>): Either<StorageFailure, List<OtherUserMinimized>>
     suspend fun getNameAndHandle(userId: UserId): Either<StorageFailure, NameAndHandle>
     suspend fun migrateUserToTeam(teamName: String): Either<CoreFailure, CreateUserTeam>
+    suspend fun updateTeamId(userId: UserId, teamId: TeamId): Either<StorageFailure, Unit>
 }
 
 @Suppress("LongParameterList", "TooManyFunctions")
@@ -652,7 +653,7 @@ internal class UserDataSource internal constructor(
 
     override suspend fun migrateUserToTeam(teamName: String): Either<CoreFailure, CreateUserTeam> {
         return wrapApiRequest { upgradePersonalToTeamApi.migrateToTeam(teamName) }.map { dto ->
-            CreateUserTeam(dto.teamName)
+            CreateUserTeam(dto.teamId, dto.teamName)
         }
             .onSuccess {
                 kaliumLogger.d("Migrated user to team")
@@ -663,6 +664,10 @@ internal class UserDataSource internal constructor(
             }
     }
 
+    override suspend fun updateTeamId(userId: UserId, teamId: TeamId): Either<StorageFailure, Unit> = wrapStorageRequest {
+        userDAO.updateTeamId(userId.toDao(), teamId.value)
+    }
+
     companion object {
         internal const val SELF_USER_ID_KEY = "selfUserID"
 

logic/src/commonMain/kotlin/com/wire/kalium/logic/feature/UserSessionScope.kt

@@ -560,6 +560,10 @@ class UserSessionScope internal constructor(
         }
     }
 
+    private val invalidateTeamId = {
+        _teamId = Either.Left(CoreFailure.Unknown(Throwable("NotInitialized")))
+    }
+
     private val selfTeamId = SelfTeamIdProvider { teamId() }
 
     private val accessTokenRepository: AccessTokenRepository
@@ -2097,7 +2101,7 @@ class UserSessionScope internal constructor(
         )
 
     val migrateFromPersonalToTeam: MigrateFromPersonalToTeamUseCase
-        get() = MigrateFromPersonalToTeamUseCaseImpl(userRepository)
+        get() = MigrateFromPersonalToTeamUseCaseImpl(userId, userRepository, invalidateTeamId)
 
     internal val getProxyCredentials: GetProxyCredentialsUseCase
         get() = GetProxyCredentialsUseCaseImpl(sessionManager)

logic/src/commonMain/kotlin/com/wire/kalium/logic/feature/user/migration/MigrateFromPersonalToTeamUseCase.kt

@@ -18,6 +18,8 @@
 package com.wire.kalium.logic.feature.user.migration
 
 import com.wire.kalium.logic.CoreFailure
+import com.wire.kalium.logic.data.id.TeamId
+import com.wire.kalium.logic.data.user.UserId
 import com.wire.kalium.logic.data.user.UserRepository
 import com.wire.kalium.logic.functional.fold
 
@@ -30,22 +32,25 @@ interface MigrateFromPersonalToTeamUseCase {
 }
 
 sealed class MigrateFromPersonalToTeamResult {
-    data class Success(val teamName: String) : MigrateFromPersonalToTeamResult()
+    data object Success : MigrateFromPersonalToTeamResult()
     data class Error(val failure: CoreFailure) : MigrateFromPersonalToTeamResult()
 }
 
 internal class MigrateFromPersonalToTeamUseCaseImpl internal constructor(
+    private val selfUserId: UserId,
     private val userRepository: UserRepository,
+    private val invalidateTeamId: () -> Unit
 ) : MigrateFromPersonalToTeamUseCase {
     override suspend operator fun invoke(
         teamName: String,
     ): MigrateFromPersonalToTeamResult {
         return userRepository.migrateUserToTeam(teamName)
             .fold(
                 { error -> return MigrateFromPersonalToTeamResult.Error(error) },
-                { success ->
-                    // TODO Invalidate team id in memory so UserSessionScope.selfTeamId got updated data WPB-12187
-                    MigrateFromPersonalToTeamResult.Success(teamName = success.teamName)
+                { user ->
+                    userRepository.updateTeamId(selfUserId, TeamId(user.teamId))
+                    invalidateTeamId()
+                    MigrateFromPersonalToTeamResult.Success
                 }
             )
     }

logic/src/commonTest/kotlin/com/wire/kalium/logic/feature/user/migration/MigrateFromPersonalToTeamUseCaseTest.kt

@@ -20,37 +20,49 @@ package com.wire.kalium.logic.feature.user.migration
 
 import com.wire.kalium.logic.CoreFailure
 import com.wire.kalium.logic.NetworkFailure
+import com.wire.kalium.logic.StorageFailure
 import com.wire.kalium.logic.data.user.CreateUserTeam
 import com.wire.kalium.logic.data.user.UserRepository
+import com.wire.kalium.logic.framework.TestUser
 import com.wire.kalium.logic.functional.Either
 import com.wire.kalium.network.api.model.ErrorResponse
 import com.wire.kalium.network.exceptions.KaliumException
 import io.ktor.http.HttpStatusCode
 import io.mockative.Mock
 import io.mockative.any
 import io.mockative.coEvery
+import io.mockative.coVerify
 import io.mockative.mock
+import io.mockative.once
 import kotlinx.coroutines.test.runTest
 import kotlin.test.Test
 import kotlin.test.assertEquals
 import kotlin.test.assertIs
+import kotlin.test.assertTrue
 
 class MigrateFromPersonalToTeamUseCaseTest {
 
     @Test
     fun givenRepositorySucceeds_whenMigratingUserToTeam_thenShouldPropagateSuccess() = runTest {
-        val (_, useCase) = Arrangement().withSuccessRepository().arrange()
+        val (arrangement, useCase) = Arrangement()
+            .withUpdateTeamIdReturning(Either.Right(Unit))
+            .withMigrationSuccess()
+            .arrange()
 
         val result = useCase(teamName = "teamName")
 
+        coVerify {
+            arrangement.userRepository.updateTeamId(any(), any())
+        }.wasInvoked(exactly = once)
+        assertTrue(arrangement.isCachedTeamIdInvalidated)
         assertIs<MigrateFromPersonalToTeamResult.Success>(result)
     }
 
     @Test
     fun givenRepositoryFailsWithNoNetworkConnection_whenMigratingUserToTeam_thenShouldPropagateFailure() =
         runTest {
             val coreFailure = NetworkFailure.NoNetworkConnection(null)
-            val (_, useCase) = Arrangement().withRepositoryReturning(Either.Left(coreFailure))
+            val (_, useCase) = Arrangement().withMigrationReturning(Either.Left(coreFailure))
                 .arrange()
 
             val result = useCase(teamName = "teamName")
@@ -79,7 +91,7 @@ class MigrateFromPersonalToTeamUseCaseTest {
     @Test
     fun givenRepositoryFailsWithNotFound_whenMigratingUserToTeam_thenShouldPropagateFailure() =
         runTest {
-            val (_, useCase) = Arrangement().withNotFoundRepository().arrange()
+            val (_, useCase) = Arrangement().withMigrationFailure().arrange()
 
             val result = useCase(teamName = "teamName")
 
@@ -98,10 +110,12 @@ class MigrateFromPersonalToTeamUseCaseTest {
         @Mock
         val userRepository: UserRepository = mock(UserRepository::class)
 
-        suspend fun withSuccessRepository() = apply {
+        var isCachedTeamIdInvalidated = false
+
+        suspend fun withMigrationSuccess() = apply {
             coEvery { userRepository.migrateUserToTeam(any()) }.returns(
                 Either.Right(
-                    CreateUserTeam("teamName")
+                    CreateUserTeam("teamId", "teamName")
                 )
             )
         }
@@ -122,7 +136,7 @@ class MigrateFromPersonalToTeamUseCaseTest {
             )
         }
 
-        suspend fun withNotFoundRepository() = apply {
+        suspend fun withMigrationFailure() = apply {
             coEvery { userRepository.migrateUserToTeam(any()) }.returns(
                 Either.Left(
                     NetworkFailure.ServerMiscommunication(
@@ -138,14 +152,21 @@ class MigrateFromPersonalToTeamUseCaseTest {
             )
         }
 
-        suspend fun withRepositoryReturning(result: Either<CoreFailure, CreateUserTeam>) =
+        suspend fun withMigrationReturning(result: Either<CoreFailure, CreateUserTeam>) =
             apply {
                 coEvery { userRepository.migrateUserToTeam(any()) }.returns(result)
             }
 
+        suspend fun withUpdateTeamIdReturning(result: Either<StorageFailure, Unit>) = apply {
+            coEvery { userRepository.updateTeamId(any(), any()) }.returns(result)
+        }
+
         fun arrange() = this to MigrateFromPersonalToTeamUseCaseImpl(
-            userRepository = userRepository
+            selfUserId = TestUser.SELF.id,
+            userRepository = userRepository,
+            invalidateTeamId = {
+                isCachedTeamIdInvalidated = true
+            }
         )
     }
-
-}
+}

persistence/src/commonMain/db_user/com/wire/kalium/persistence/Users.sq

@@ -273,3 +273,6 @@ SELECT active_one_on_one_conversation_id FROM User WHERE qualified_id = :userId;
 
 selectNamesAndHandle:
 SELECT name, handle FROM User WHERE qualified_id = :userId;
+
+updateTeamId:
+UPDATE User SET team = ? WHERE qualified_id = ?;

persistence/src/commonMain/kotlin/com/wire/kalium/persistence/dao/UserDAO.kt

@@ -316,4 +316,5 @@ interface UserDAO {
     suspend fun getOneOnOnConversationId(userId: UserIDEntity): QualifiedIDEntity?
     suspend fun getUsersMinimizedByQualifiedIDs(qualifiedIDs: List<QualifiedIDEntity>): List<UserEntityMinimized>
     suspend fun getNameAndHandle(userId: UserIDEntity): NameAndHandleEntity?
+    suspend fun updateTeamId(userId: UserIDEntity, teamId: String)
 }

persistence/src/commonMain/kotlin/com/wire/kalium/persistence/dao/UserDAOImpl.kt

@@ -473,4 +473,8 @@ class UserDAOImpl internal constructor(
     override suspend fun getNameAndHandle(userId: UserIDEntity): NameAndHandleEntity? = withContext(queriesContext) {
         userQueries.selectNamesAndHandle(userId, ::NameAndHandleEntity).executeAsOneOrNull()
     }
+
+    override suspend fun updateTeamId(userId: UserIDEntity, teamId: String) {
+        userQueries.updateTeamId(teamId, userId)
+    }
 }