Fix message backup checkout e2e tests.

app/src/androidTest/java/org/thoughtcrime/securesms/backup/v2/ui/subscription/MessageBackupsCheckoutActivityTest.kt

@@ -20,9 +20,12 @@ import assertk.assertions.isEqualTo
 import assertk.assertions.isNull
 import io.mockk.coEvery
 import io.mockk.every
+import io.mockk.mockk
+import io.mockk.mockkObject
 import io.mockk.mockkStatic
 import kotlinx.coroutines.flow.MutableSharedFlow
 import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.test.StandardTestDispatcher
 import org.junit.Before
 import org.junit.Rule
 import org.junit.Test
@@ -38,9 +41,12 @@ import org.thoughtcrime.securesms.database.InAppPaymentTable
 import org.thoughtcrime.securesms.database.SignalDatabase
 import org.thoughtcrime.securesms.dependencies.AppDependencies
 import org.thoughtcrime.securesms.keyvalue.SignalStore
+import org.thoughtcrime.securesms.net.SignalNetwork
+import org.thoughtcrime.securesms.testing.CoroutineDispatcherRule
 import org.thoughtcrime.securesms.testing.InAppPaymentsRule
 import org.thoughtcrime.securesms.testing.SignalActivityRule
 import org.thoughtcrime.securesms.util.RemoteConfig
+import org.whispersystems.signalservice.api.NetworkResult
 import java.math.BigDecimal
 import java.util.Currency
 
@@ -53,6 +59,10 @@ class MessageBackupsCheckoutActivityTest {
 
   @get:Rule val composeTestRule = createEmptyComposeRule()
 
+  private val testDispatcher = StandardTestDispatcher()
+
+  @get:Rule val coroutineDispatcherRule = CoroutineDispatcherRule(testDispatcher)
+
   private val purchaseResults = MutableSharedFlow<BillingPurchaseResult>()
 
   @Before
@@ -61,6 +71,11 @@ class MessageBackupsCheckoutActivityTest {
     coEvery { AppDependencies.billingApi.queryProduct() } returns BillingProduct(price = FiatMoney(BigDecimal.ONE, Currency.getInstance("USD")))
     coEvery { AppDependencies.billingApi.launchBillingFlow(any()) } returns Unit
 
+    mockkObject(SignalNetwork)
+    every { SignalNetwork.archive } returns mockk {
+      every { triggerBackupIdReservation(any(), any(), any()) } returns NetworkResult.Success(Unit)
+    }
+
     mockkStatic(RemoteConfig::class)
     every { RemoteConfig.messageBackups } returns true
   }
@@ -79,6 +94,8 @@ class MessageBackupsCheckoutActivityTest {
     composeTestRule.onNodeWithText(context.getString(R.string.MessageBackupsTypeSelectionScreen__next)).performClick()
     composeTestRule.waitForIdle()
 
+    testDispatcher.scheduler.advanceUntilIdle()
+
     runBlocking {
       purchaseResults.emit(
         BillingPurchaseResult.Success(
@@ -94,6 +111,8 @@ class MessageBackupsCheckoutActivityTest {
     composeTestRule.waitForIdle()
     composeTestRule.onNodeWithTag("dialog-circular-progress-indicator").assertIsDisplayed()
 
+    testDispatcher.scheduler.advanceUntilIdle()
+
     val iap = SignalDatabase.inAppPayments.getLatestInAppPaymentByType(InAppPaymentType.RECURRING_BACKUP)
     assertThat(iap?.state).isEqualTo(InAppPaymentTable.State.PENDING)
 
@@ -162,8 +181,11 @@ class MessageBackupsCheckoutActivityTest {
   }
 
   private fun launchCheckoutFlow(tier: MessageBackupTier? = null): ActivityScenario<MessageBackupsCheckoutActivity> {
-    return ActivityScenario.launch(
+    val scenario = ActivityScenario.launch<MessageBackupsCheckoutActivity>(
       MessageBackupsCheckoutActivity.Contract().createIntent(InstrumentationRegistry.getInstrumentation().targetContext, tier)
     )
+
+    testDispatcher.scheduler.advanceUntilIdle()
+    return scenario
   }
 }

app/src/androidTest/java/org/thoughtcrime/securesms/testing/CoroutineDispatcherRule.kt

@@ -0,0 +1,44 @@
+/*
+ * Copyright 2025 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.thoughtcrime.securesms.testing
+
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.test.TestDispatcher
+import org.junit.rules.ExternalResource
+import org.signal.core.util.concurrent.SignalDispatchers
+
+/**
+ * Rule that allows for injection of test dispatchers when operating with ViewModels.
+ */
+class CoroutineDispatcherRule(
+  defaultDispatcher: TestDispatcher,
+  mainDispatcher: TestDispatcher = defaultDispatcher,
+  ioDispatcher: TestDispatcher = defaultDispatcher,
+  unconfinedDispatcher: TestDispatcher = defaultDispatcher
+) : ExternalResource() {
+
+  private val testDispatcherProvider = TestDispatcherProvider(
+    main = mainDispatcher,
+    io = ioDispatcher,
+    default = defaultDispatcher,
+    unconfined = unconfinedDispatcher
+  )
+
+  override fun before() {
+    SignalDispatchers.setDispatcherProvider(testDispatcherProvider)
+  }
+
+  override fun after() {
+    SignalDispatchers.setDispatcherProvider()
+  }
+
+  private class TestDispatcherProvider(
+    override val main: CoroutineDispatcher,
+    override val io: CoroutineDispatcher,
+    override val default: CoroutineDispatcher,
+    override val unconfined: CoroutineDispatcher
+  ) : SignalDispatchers.DispatcherProvider
+}

app/src/main/java/org/thoughtcrime/securesms/backup/v2/BackupRepository.kt

@@ -10,8 +10,6 @@ import android.os.Environment
 import android.os.StatFs
 import androidx.annotation.Discouraged
 import androidx.annotation.WorkerThread
-import kotlinx.coroutines.Dispatchers
-import kotlinx.coroutines.withContext
 import okio.ByteString
 import okio.ByteString.Companion.toByteString
 import org.greenrobot.eventbus.EventBus
@@ -305,9 +303,7 @@ object BackupRepository {
     }
 
     val paidType = try {
-      withContext(Dispatchers.IO) {
-        getPaidType()
-      }
+      getPaidType()
     } catch (e: IOException) {
       Log.w(TAG, "Failed to retrieve paid type.", e)
       return false
@@ -1415,7 +1411,8 @@ object BackupRepository {
     }
   }
 
-  private suspend fun getFreeType(): MessageBackupsType.Free {
+  @WorkerThread
+  private fun getFreeType(): MessageBackupsType.Free {
     val config = getSubscriptionsConfiguration()
 
     return MessageBackupsType.Free(
@@ -1426,6 +1423,7 @@ object BackupRepository {
   private suspend fun getPaidType(): MessageBackupsType.Paid? {
     val config = getSubscriptionsConfiguration()
     val product = AppDependencies.billingApi.queryProduct() ?: return null
+
     val backupLevelConfiguration = config.backupConfiguration.backupLevelConfigurationMap[SubscriptionsConfiguration.BACKUPS_LEVEL] ?: return null
 
     return MessageBackupsType.Paid(
@@ -1435,12 +1433,11 @@ object BackupRepository {
     )
   }
 
-  private suspend fun getSubscriptionsConfiguration(): SubscriptionsConfiguration {
-    val serviceResponse = withContext(Dispatchers.IO) {
-      AppDependencies
-        .donationsService
-        .getDonationsConfiguration(Locale.getDefault())
-    }
+  @WorkerThread
+  private fun getSubscriptionsConfiguration(): SubscriptionsConfiguration {
+    val serviceResponse = AppDependencies
+      .donationsService
+      .getDonationsConfiguration(Locale.getDefault())
 
     if (serviceResponse.result.isEmpty) {
       if (serviceResponse.applicationError.isPresent) {

app/src/main/java/org/thoughtcrime/securesms/backup/v2/ui/BackupAlertDelegate.kt

@@ -9,7 +9,9 @@ import androidx.fragment.app.FragmentManager
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.coroutineScope
 import androidx.lifecycle.repeatOnLifecycle
+import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.withContext
 import org.thoughtcrime.securesms.backup.v2.BackupRepository
 import org.thoughtcrime.securesms.keyvalue.SignalStore
 
@@ -25,7 +27,7 @@ object BackupAlertDelegate {
           BackupAlertBottomSheet.create(BackupAlert.BackupFailed).show(fragmentManager, null)
         } else if (BackupRepository.shouldDisplayCouldNotCompleteBackupSheet()) {
           BackupAlertBottomSheet.create(BackupAlert.CouldNotCompleteBackup(daysSinceLastBackup = SignalStore.backup.daysSinceLastBackup)).show(fragmentManager, null)
-        } else if (BackupRepository.shouldDisplayYourMediaWillBeDeletedTodaySheet()) {
+        } else if (withContext(Dispatchers.IO) { BackupRepository.shouldDisplayYourMediaWillBeDeletedTodaySheet() }) {
           BackupAlertBottomSheet.create(BackupAlert.MediaWillBeDeletedToday).show(fragmentManager, null)
         }
       }

app/src/main/java/org/thoughtcrime/securesms/backup/v2/ui/subscription/MessageBackupsFlowViewModel.kt

@@ -8,7 +8,6 @@ package org.thoughtcrime.securesms.backup.v2.ui.subscription
 import androidx.annotation.WorkerThread
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
-import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.FlowPreview
 import kotlinx.coroutines.TimeoutCancellationException
 import kotlinx.coroutines.flow.MutableStateFlow
@@ -23,6 +22,7 @@ import kotlinx.coroutines.launch
 import kotlinx.coroutines.reactive.asFlow
 import kotlinx.coroutines.withContext
 import org.signal.core.util.billing.BillingPurchaseResult
+import org.signal.core.util.concurrent.SignalDispatchers
 import org.signal.core.util.logging.Log
 import org.signal.donations.InAppPaymentType
 import org.signal.donations.PaymentSourceType
@@ -70,7 +70,7 @@ class MessageBackupsFlowViewModel(
     check(SignalStore.backup.backupTier != MessageBackupTier.PAID) { "This screen does not support cancellation or downgrades." }
 
     viewModelScope.launch {
-      val result = withContext(Dispatchers.IO) {
+      val result = withContext(SignalDispatchers.IO) {
         BackupRepository.triggerBackupIdReservation()
       }
 
@@ -79,20 +79,22 @@ class MessageBackupsFlowViewModel(
         internalStateFlow.update { it.copy(paymentReadyState = MessageBackupsFlowState.PaymentReadyState.READY) }
       }
 
-      result.runOnStatusCodeError {
-        Log.d(TAG, "Failed to trigger backup id reservation. ($it)")
+      result.runOnStatusCodeError { code ->
+        Log.d(TAG, "Failed to trigger backup id reservation. ($code)")
         internalStateFlow.update { it.copy(paymentReadyState = MessageBackupsFlowState.PaymentReadyState.FAILED) }
       }
     }
 
     viewModelScope.launch {
-      internalStateFlow.update {
-        it.copy(
-          availableBackupTypes = BackupRepository.getAvailableBackupsTypes(
-            if (!RemoteConfig.messageBackups) emptyList() else listOf(MessageBackupTier.FREE, MessageBackupTier.PAID)
-          )
+      val availableBackupTypes = withContext(SignalDispatchers.IO) {
+        BackupRepository.getAvailableBackupsTypes(
+          if (!RemoteConfig.messageBackups) emptyList() else listOf(MessageBackupTier.FREE, MessageBackupTier.PAID)
         )
       }
+
+      internalStateFlow.update {
+        it.copy(availableBackupTypes = availableBackupTypes)
+      }
     }
 
     viewModelScope.launch {
@@ -218,7 +220,7 @@ class MessageBackupsFlowViewModel(
         check(state.selectedMessageBackupTier == MessageBackupTier.PAID)
         check(state.availableBackupTypes.any { it.tier == state.selectedMessageBackupTier })
 
-        viewModelScope.launch(Dispatchers.IO) {
+        viewModelScope.launch(SignalDispatchers.IO) {
           internalStateFlow.update { it.copy(inAppPayment = null) }
 
           val paidFiat = AppDependencies.billingApi.queryProduct()!!.price
@@ -265,7 +267,7 @@ class MessageBackupsFlowViewModel(
    */
   @OptIn(FlowPreview::class)
   private suspend fun handleSuccess(result: BillingPurchaseResult.Success, inAppPaymentId: InAppPaymentTable.InAppPaymentId) {
-    withContext(Dispatchers.IO) {
+    withContext(SignalDispatchers.IO) {
       Log.d(TAG, "Setting purchase token data on InAppPayment and InAppPaymentSubscriber.")
       ensureSubscriberIdForBackups(IAPSubscriptionId.GooglePlayBillingPurchaseToken(result.purchaseToken))
 
@@ -287,7 +289,7 @@ class MessageBackupsFlowViewModel(
       InAppPaymentPurchaseTokenJob.createJobChain(inAppPayment).enqueue()
     }
 
-    val terminalInAppPayment = withContext(Dispatchers.IO) {
+    val terminalInAppPayment = withContext(SignalDispatchers.IO) {
       Log.d(TAG, "Awaiting completion of job chain for up to 10 seconds.")
       InAppPaymentsRepository.observeUpdates(inAppPaymentId).asFlow()
         .filter { it.state == InAppPaymentTable.State.END }

core-util/src/main/java/org/signal/core/util/concurrent/SignalDispatchers.kt

@@ -0,0 +1,40 @@
+/*
+ * Copyright 2025 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.signal.core.util.concurrent
+
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.Dispatchers
+
+/**
+ * [Dispatchers] wrapper to allow tests to inject test dispatchers.
+ */
+object SignalDispatchers {
+
+  private var dispatcherProvider: DispatcherProvider = DefaultDispatcherProvider
+
+  fun setDispatcherProvider(dispatcherProvider: DispatcherProvider = DefaultDispatcherProvider) {
+    this.dispatcherProvider = dispatcherProvider
+  }
+
+  val Main get() = dispatcherProvider.main
+  val IO get() = dispatcherProvider.io
+  val Default get() = dispatcherProvider.default
+  val Unconfined get() = dispatcherProvider.unconfined
+
+  interface DispatcherProvider {
+    val main: CoroutineDispatcher
+    val io: CoroutineDispatcher
+    val default: CoroutineDispatcher
+    val unconfined: CoroutineDispatcher
+  }
+
+  private object DefaultDispatcherProvider : DispatcherProvider {
+    override val main: CoroutineDispatcher = Dispatchers.Main
+    override val io: CoroutineDispatcher = Dispatchers.IO
+    override val default: CoroutineDispatcher = Dispatchers.Default
+    override val unconfined: CoroutineDispatcher = Dispatchers.Unconfined
+  }
+}