From 8b1c2bf2af0c9584709e5482d9e30701bf23bc58 Mon Sep 17 00:00:00 2001
From: Greyson Parrelli <greyson@signal.org>
Date: Wed, 22 Jul 2026 12:14:46 -0400
Subject: [PATCH] Fix RegistrationEndToEndTest flakiness.

---
 .../registration/RegistrationEndToEndTest.kt  | 12 ++++++---
 .../fakes/FakeStorageController.kt            | 25 +++++++++++++------
 2 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/feature/registration/src/test/java/org/signal/registration/RegistrationEndToEndTest.kt b/feature/registration/src/test/java/org/signal/registration/RegistrationEndToEndTest.kt
index 11b73af47a..d06a0d48b5 100644
--- a/feature/registration/src/test/java/org/signal/registration/RegistrationEndToEndTest.kt
+++ b/feature/registration/src/test/java/org/signal/registration/RegistrationEndToEndTest.kt
@@ -118,7 +118,6 @@ class RegistrationEndToEndTest {
     networkController = FakeNetworkController()
     storageController = FakeStorageController()
     repository = RegistrationRepository(context, networkController, storageController, isLinkAndSyncAvailable = false)
-    viewModel = RegistrationViewModel(repository, SavedStateHandle())
   }
 
   @Test
@@ -952,10 +951,14 @@ class RegistrationEndToEndTest {
     enterPhoneNumber()
     enterAep(wrongAep)
 
-    // The server rejects the recovery password derived from the wrong AEP, disabling submission until the key changes
-    waitFor("the incorrect AEP to be rejected") {
+    // The server rejects the recovery password derived from the wrong AEP, disabling submission until the key changes.
+    // Wait on the error text rather than the disabled button, which is also disabled while the attempt is in flight.
+    waitForText(ApplicationProvider.getApplicationContext<Application>().getString(R.string.EnterAepScreen__incorrect_recovery_key))
+    assert(
       composeTestRule.onAllNodesWithTag(TestTags.ENTER_AEP_NEXT_BUTTON).fetchSemanticsNodes().firstOrNull()
         ?.config?.getOrNull(SemanticsProperties.Disabled) != null
+    ) {
+      "Expected submission to be disabled after the AEP was rejected"
     }
     assert(networkController.lastRegisterAccountRequest?.recoveryPassword == wrongAep.deriveMasterKey().deriveRegistrationRecoveryPassword()) {
       "Expected a registration attempt with the wrong recovery password but was ${networkController.lastRegisterAccountRequest}"
@@ -1295,6 +1298,9 @@ class RegistrationEndToEndTest {
     folderPickerResult: Uri? = null,
     onRegistrationComplete: () -> Unit = {}
   ) {
+    // Created here rather than in setup() so that its async init cannot read controller state before the test has finished configuring it
+    viewModel = RegistrationViewModel(repository, SavedStateHandle())
+
     composeTestRule.setContent {
       SignalTheme {
         ActivityResultInterceptor(folderPickerResult) {
diff --git a/feature/registration/src/test/java/org/signal/registration/fakes/FakeStorageController.kt b/feature/registration/src/test/java/org/signal/registration/fakes/FakeStorageController.kt
index e8d25762da..53c24eb029 100644
--- a/feature/registration/src/test/java/org/signal/registration/fakes/FakeStorageController.kt
+++ b/feature/registration/src/test/java/org/signal/registration/fakes/FakeStorageController.kt
@@ -29,6 +29,9 @@ import java.time.LocalDateTime
  */
 class FakeStorageController : StorageController {
 
+  /** Guards [inProgressData], whose read-modify-write updates arrive concurrently from multiple [kotlinx.coroutines.Dispatchers.IO] threads. */
+  private val dataLock = Any()
+
   var inProgressData: RegistrationData = RegistrationData()
     private set
   var committedData: RegistrationData? = null
@@ -73,24 +76,30 @@ class FakeStorageController : StorageController {
   override suspend fun getPreExistingRegistrationData(): PreExistingRegistrationData? = preExistingRegistrationData
 
   override suspend fun clearAllData() {
-    inProgressData = RegistrationData()
+    synchronized(dataLock) {
+      inProgressData = RegistrationData()
+    }
   }
 
   override suspend fun clearLocalDataAndRestart() = notExpected()
 
-  override suspend fun readInProgressRegistrationData(): RegistrationData = inProgressData
+  override suspend fun readInProgressRegistrationData(): RegistrationData = synchronized(dataLock) { inProgressData }
 
   override suspend fun updateInProgressRegistrationData(updater: RegistrationData.Builder.() -> Unit) {
-    inProgressData = inProgressData.newBuilder().apply(updater).lastUpdatedMillis(System.currentTimeMillis()).build()
+    synchronized(dataLock) {
+      inProgressData = inProgressData.newBuilder().apply(updater).lastUpdatedMillis(System.currentTimeMillis()).build()
+    }
   }
 
   override suspend fun commitRegistrationData() {
-    val accountData = inProgressData.accountData
-    val accountDataComplete = accountData != null && accountData.e164.isNotEmpty() && accountData.aci.isNotEmpty() && accountData.pni.isNotEmpty() && accountData.servicePassword.isNotEmpty()
-    if (!inProgressData.accountDataCommitted && accountDataComplete) {
-      inProgressData = inProgressData.newBuilder().accountDataCommitted(true).build()
+    synchronized(dataLock) {
+      val accountData = inProgressData.accountData
+      val accountDataComplete = accountData != null && accountData.e164.isNotEmpty() && accountData.aci.isNotEmpty() && accountData.pni.isNotEmpty() && accountData.servicePassword.isNotEmpty()
+      if (!inProgressData.accountDataCommitted && accountDataComplete) {
+        inProgressData = inProgressData.newBuilder().accountDataCommitted(true).build()
+      }
+      committedData = inProgressData
     }
-    committedData = inProgressData
   }
 
   override suspend fun setRestoreDecision(decision: RestoreDecision) {