Wait for UnifiedPush pusher to be registered before returning Result.

libraries/push/impl/src/main/kotlin/io/element/android/libraries/push/impl/DefaultPushService.kt

@@ -67,13 +67,13 @@ class DefaultPushService @Inject constructor(
             pushProviders.find { it.name == currentPushProviderName }?.unregister(matrixClient)
                 ?.onFailure {
                     Timber.w(it, "Failed to unregister previous push provider")
+                    return Result.failure(it)
                 }
         }
+        // Store new value
+        userPushStore.setPushProviderName(pushProvider.name)
+        // Then try to register
         return pushProvider.registerWith(matrixClient, distributor)
-            .onSuccess {
-                // Store new value
-                userPushStore.setPushProviderName(pushProvider.name)
-            }
     }
 
     override suspend fun testPush(): Boolean {

libraries/pushproviders/unifiedpush/src/main/kotlin/io/element/android/libraries/pushproviders/unifiedpush/RegisterUnifiedPushUseCase.kt

@@ -19,17 +19,41 @@ package io.element.android.libraries.pushproviders.unifiedpush
 import android.content.Context
 import io.element.android.libraries.di.ApplicationContext
 import io.element.android.libraries.pushproviders.api.Distributor
+import io.element.android.libraries.pushproviders.unifiedpush.registration.EndpointRegistrationHandler
+import kotlinx.coroutines.CompletableDeferred
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.flow.filter
+import kotlinx.coroutines.flow.first
+import kotlinx.coroutines.launch
+import kotlinx.coroutines.withTimeout
 import org.unifiedpush.android.connector.UnifiedPush
 import javax.inject.Inject
+import kotlin.time.Duration.Companion.seconds
 
 class RegisterUnifiedPushUseCase @Inject constructor(
     @ApplicationContext private val context: Context,
+    private val endpointRegistrationHandler: EndpointRegistrationHandler,
+    private val coroutineScope: CoroutineScope,
 ) {
-    fun execute(distributor: Distributor, clientSecret: String): Result<Unit> {
+    suspend fun execute(distributor: Distributor, clientSecret: String): Result<Unit> {
         UnifiedPush.saveDistributor(context, distributor.value)
+        val completable = CompletableDeferred<Result<Unit>>()
+        val job = coroutineScope.launch {
+            val result = endpointRegistrationHandler.state
+                .filter { it.clientSecret == clientSecret }
+                .first()
+                .result
+            completable.complete(result)
+        }
         // This will trigger the callback
         // VectorUnifiedPushMessagingReceiver.onNewEndpoint
         UnifiedPush.registerApp(context = context, instance = clientSecret)
-        return Result.success(Unit)
+        // Wait for VectorUnifiedPushMessagingReceiver.onNewEndpoint to proceed
+        return withTimeout(30.seconds) {
+            completable.await()
+        }
+            .onFailure {
+                job.cancel()
+            }
     }
 }

libraries/pushproviders/unifiedpush/src/main/kotlin/io/element/android/libraries/pushproviders/unifiedpush/VectorUnifiedPushMessagingReceiver.kt

@@ -21,6 +21,8 @@ import android.content.Intent
 import io.element.android.libraries.architecture.bindings
 import io.element.android.libraries.core.log.logger.LoggerTag
 import io.element.android.libraries.pushproviders.api.PushHandler
+import io.element.android.libraries.pushproviders.unifiedpush.registration.EndpointRegistrationHandler
+import io.element.android.libraries.pushproviders.unifiedpush.registration.RegistrationResult
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.launch
@@ -37,6 +39,7 @@ class VectorUnifiedPushMessagingReceiver : MessagingReceiver() {
     @Inject lateinit var unifiedPushStore: UnifiedPushStore
     @Inject lateinit var unifiedPushGatewayResolver: UnifiedPushGatewayResolver
     @Inject lateinit var newGatewayHandler: UnifiedPushNewGatewayHandler
+    @Inject lateinit var endpointRegistrationHandler: EndpointRegistrationHandler
 
     private val coroutineScope = CoroutineScope(SupervisorJob())
 
@@ -73,14 +76,28 @@ class VectorUnifiedPushMessagingReceiver : MessagingReceiver() {
         coroutineScope.launch {
             val gateway = unifiedPushGatewayResolver.getGateway(endpoint)
             unifiedPushStore.storePushGateway(gateway, instance)
-            gateway?.let { pushGateway ->
-                newGatewayHandler.handle(endpoint, pushGateway, instance)
+            if (gateway == null) {
+                Timber.tag(loggerTag.value).w("No gateway found for endpoint $endpoint")
+                endpointRegistrationHandler.registrationDone(
+                    RegistrationResult(
+                        clientSecret = instance,
+                        result = Result.failure(IllegalStateException("No gateway found for endpoint $endpoint")),
+                    )
+                )
+            } else {
+                val result = newGatewayHandler.handle(endpoint, gateway, instance)
                     .onFailure {
                         Timber.tag(loggerTag.value).e(it, "Failed to handle new gateway")
                     }
                     .onSuccess {
                         unifiedPushStore.storeUpEndpoint(endpoint, instance)
                     }
+                endpointRegistrationHandler.registrationDone(
+                    RegistrationResult(
+                        clientSecret = instance,
+                        result = result,
+                    )
+                )
             }
         }
         guardServiceStarter.stop()

libraries/pushproviders/unifiedpush/src/main/kotlin/io/element/android/libraries/pushproviders/unifiedpush/registration/EndpointRegistrationHandler.kt

@@ -0,0 +1,38 @@
+/*
+ * Copyright (c) 2024 New Vector Ltd
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package io.element.android.libraries.pushproviders.unifiedpush.registration
+
+import io.element.android.libraries.di.AppScope
+import io.element.android.libraries.di.SingleIn
+import kotlinx.coroutines.flow.MutableSharedFlow
+import kotlinx.coroutines.flow.SharedFlow
+import javax.inject.Inject
+
+data class RegistrationResult(
+    val clientSecret: String,
+    val result: Result<Unit>,
+)
+
+@SingleIn(AppScope::class)
+class EndpointRegistrationHandler @Inject constructor() {
+    private val _state = MutableSharedFlow<RegistrationResult>()
+    val state: SharedFlow<RegistrationResult> = _state
+
+    suspend fun registrationDone(result: RegistrationResult) {
+        _state.emit(result)
+    }
+}