Secure Storage: improve API for waitForBackupUploadSteadyState()

features/logout/impl/src/main/kotlin/io/element/android/features/logout/impl/LogoutPresenter.kt

@@ -29,6 +29,7 @@ import io.element.android.libraries.architecture.Async
 import io.element.android.libraries.architecture.Presenter
 import io.element.android.libraries.architecture.runCatchingUpdatingState
 import io.element.android.libraries.matrix.api.MatrixClient
+import io.element.android.libraries.matrix.api.encryption.BackupUploadState
 import io.element.android.libraries.matrix.api.encryption.EncryptionService
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.launch
@@ -46,13 +47,15 @@ class LogoutPresenter @Inject constructor(
             mutableStateOf(Async.Uninitialized)
         }
 
-        val backupUploadState by encryptionService.backupUploadStateStateFlow.collectAsState()
+        val backupUploadState: BackupUploadState by remember {
+            encryptionService.waitForBackupUploadSteadyState()
+        }
+            .collectAsState(initial = BackupUploadState.Unknown)
 
         var showLogoutDialog by remember { mutableStateOf(false) }
         var isLastSession by remember { mutableStateOf(false) }
         LaunchedEffect(Unit) {
             isLastSession = encryptionService.isLastDevice().getOrNull() ?: false
-            encryptionService.waitForBackupUploadSteadyState()
         }
 
         fun handleEvents(event: LogoutEvents) {

features/logout/impl/src/test/kotlin/io/element/android/features/logout/impl/LogoutPresenterTest.kt

@@ -28,6 +28,8 @@ import io.element.android.libraries.matrix.test.A_THROWABLE
 import io.element.android.libraries.matrix.test.FakeMatrixClient
 import io.element.android.libraries.matrix.test.encryption.FakeEncryptionService
 import io.element.android.tests.testutils.WarmUpRule
+import kotlinx.coroutines.delay
+import kotlinx.coroutines.flow.flow
 import kotlinx.coroutines.test.runTest
 import org.junit.Rule
 import org.junit.Test
@@ -73,6 +75,15 @@ class LogoutPresenterTest {
     @Test
     fun `present - initial state - backing up`() = runTest {
         val encryptionService = FakeEncryptionService()
+        encryptionService.givenWaitForBackupUploadSteadyStateFlow(
+            flow {
+                emit(BackupUploadState.Waiting)
+                delay(1)
+                emit(BackupUploadState.Uploading(backedUpCount = 1, totalCount = 2))
+                delay(1)
+                emit(BackupUploadState.Done)
+            }
+        )
         val presenter = createLogoutPresenter(
             encryptionService = encryptionService
         )
@@ -81,13 +92,11 @@ class LogoutPresenterTest {
         }.test {
             val initialState = awaitItem()
             assertThat(initialState.isLastSession).isFalse()
-            assertThat(initialState.backupUploadState).isEqualTo(BackupUploadState.Unknown)
+            assertThat(initialState.backupUploadState).isEqualTo(BackupUploadState.Waiting)
             assertThat(initialState.showConfirmationDialog).isFalse()
             assertThat(initialState.logoutAction).isEqualTo(Async.Uninitialized)
-            encryptionService.emitBackupUploadState(BackupUploadState.Uploading(backedUpCount = 1, totalCount = 2))
             val state = awaitItem()
             assertThat(state.backupUploadState).isEqualTo(BackupUploadState.Uploading(backedUpCount = 1, totalCount = 2))
-            encryptionService.emitBackupUploadState(BackupUploadState.Done)
             val doneState = awaitItem()
             assertThat(doneState.backupUploadState).isEqualTo(BackupUploadState.Done)
         }

libraries/matrix/api/src/main/kotlin/io/element/android/libraries/matrix/api/encryption/EncryptionService.kt

@@ -16,12 +16,12 @@
 
 package io.element.android.libraries.matrix.api.encryption
 
+import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.StateFlow
 
 interface EncryptionService {
     val backupStateStateFlow: StateFlow<BackupState>
     val recoveryStateStateFlow: StateFlow<RecoveryState>
-    val backupUploadStateStateFlow: StateFlow<BackupUploadState>
     val enableRecoveryProgressStateFlow: StateFlow<EnableRecoveryProgress>
 
     suspend fun enableBackups(): Result<Unit>
@@ -46,7 +46,7 @@ interface EncryptionService {
     suspend fun fixRecoveryIssues(recoveryKey: String): Result<Unit>
 
     /**
-     * Observe [backupUploadStateStateFlow] to get progress.
+     * Wait for backup upload steady state.
      */
-    suspend fun waitForBackupUploadSteadyState(): Result<Unit>
+    fun waitForBackupUploadSteadyState(): Flow<BackupUploadState>
 }

libraries/matrix/impl/src/main/kotlin/io/element/android/libraries/matrix/impl/encryption/RustEncryptionService.kt

@@ -22,7 +22,10 @@ import io.element.android.libraries.matrix.api.encryption.BackupUploadState
 import io.element.android.libraries.matrix.api.encryption.EnableRecoveryProgress
 import io.element.android.libraries.matrix.api.encryption.EncryptionService
 import io.element.android.libraries.matrix.api.encryption.RecoveryState
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.callbackFlow
 import kotlinx.coroutines.withContext
 import org.matrix.rustcomponents.sdk.BackupStateListener
 import org.matrix.rustcomponents.sdk.BackupSteadyStateListener
@@ -50,7 +53,6 @@ internal class RustEncryptionService(
     override val backupStateStateFlow: MutableStateFlow<BackupState> = MutableStateFlow(service.backupState().let(backupStateMapper::map))
     override val recoveryStateStateFlow: MutableStateFlow<RecoveryState> = MutableStateFlow(service.recoveryState().let(recoveryStateMapper::map))
     override val enableRecoveryProgressStateFlow: MutableStateFlow<EnableRecoveryProgress> = MutableStateFlow(EnableRecoveryProgress.Unknown)
-    override val backupUploadStateStateFlow: MutableStateFlow<BackupUploadState> = MutableStateFlow(BackupUploadState.Unknown)
 
     fun start() {
         service.backupStateListener(object : BackupStateListener {
@@ -94,16 +96,19 @@ internal class RustEncryptionService(
         }
     }
 
-    override suspend fun waitForBackupUploadSteadyState(
+    override fun waitForBackupUploadSteadyState(): Flow<BackupUploadState> {
-    ): Result<Unit> = withContext(dispatchers.io) {
+        return callbackFlow {
-        runCatching {
             service.waitForBackupUploadSteadyState(
                 progressListener = object : BackupSteadyStateListener {
                     override fun onUpdate(status: RustBackupUploadState) {
-                        backupUploadStateStateFlow.value = backupUploadStateMapper.map(status)
+                        trySend(backupUploadStateMapper.map(status))
+                        if (status == RustBackupUploadState.Done) {
+                            close()
+                        }
                     }
                 }
             )
+            awaitClose {}
         }
     }
 

libraries/matrix/test/src/main/kotlin/io/element/android/libraries/matrix/test/encryption/FakeEncryptionService.kt

@@ -22,14 +22,16 @@ import io.element.android.libraries.matrix.api.encryption.EnableRecoveryProgress
 import io.element.android.libraries.matrix.api.encryption.EncryptionService
 import io.element.android.libraries.matrix.api.encryption.RecoveryState
 import io.element.android.tests.testutils.simulateLongTask
+import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.flowOf
 
 class FakeEncryptionService : EncryptionService {
     private var disableRecoveryFailure: Exception? = null
     override val backupStateStateFlow: MutableStateFlow<BackupState> = MutableStateFlow(BackupState.UNKNOWN)
     override val recoveryStateStateFlow: MutableStateFlow<RecoveryState> = MutableStateFlow(RecoveryState.UNKNOWN)
     override val enableRecoveryProgressStateFlow: MutableStateFlow<EnableRecoveryProgress> = MutableStateFlow(EnableRecoveryProgress.Unknown)
-    override val backupUploadStateStateFlow: MutableStateFlow<BackupUploadState> = MutableStateFlow(BackupUploadState.Unknown)
+    private var waitForBackupUploadSteadyStateFlow: Flow<BackupUploadState> = flowOf()
 
     private var fixRecoveryIssuesFailure: Exception? = null
 
@@ -73,12 +75,12 @@ class FakeEncryptionService : EncryptionService {
         return Result.success(Unit)
     }
 
-    override suspend fun waitForBackupUploadSteadyState(): Result<Unit> {
+    fun givenWaitForBackupUploadSteadyStateFlow(flow: Flow<BackupUploadState>) {
-        return Result.success(Unit)
+        waitForBackupUploadSteadyStateFlow = flow
     }
 
-    suspend fun emitBackupUploadState(state: BackupUploadState) {
+    override fun waitForBackupUploadSteadyState(): Flow<BackupUploadState> {
-        backupUploadStateStateFlow.emit(state)
+        return waitForBackupUploadSteadyStateFlow
     }
 
     suspend fun emitBackupState(state: BackupState) {