From 0639bc6a5b01d7a522e0a996347b3a7e8b9c22d4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jorge=20Mart=C3=ADn?= <jorgem@element.io>
Date: Thu, 23 May 2024 09:44:05 +0200
Subject: [PATCH] Simplify `RegisterUnifiedPushUseCase`, fix test

---
 .../unifiedpush/RegisterUnifiedPushUseCase.kt | 25 ++++++-------------
 .../DefaultRegisterUnifiedPushUseCaseTest.kt  |  9 ++++---
 2 files changed, 13 insertions(+), 21 deletions(-)

diff --git a/libraries/pushproviders/unifiedpush/src/main/kotlin/io/element/android/libraries/pushproviders/unifiedpush/RegisterUnifiedPushUseCase.kt b/libraries/pushproviders/unifiedpush/src/main/kotlin/io/element/android/libraries/pushproviders/unifiedpush/RegisterUnifiedPushUseCase.kt
index 1839bb4d18..6670f18ea2 100644
--- a/libraries/pushproviders/unifiedpush/src/main/kotlin/io/element/android/libraries/pushproviders/unifiedpush/RegisterUnifiedPushUseCase.kt
+++ b/libraries/pushproviders/unifiedpush/src/main/kotlin/io/element/android/libraries/pushproviders/unifiedpush/RegisterUnifiedPushUseCase.kt
@@ -22,11 +22,8 @@ import io.element.android.libraries.di.AppScope
 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
@@ -40,27 +37,21 @@ interface RegisterUnifiedPushUseCase {
 class DefaultRegisterUnifiedPushUseCase @Inject constructor(
     @ApplicationContext private val context: Context,
     private val endpointRegistrationHandler: EndpointRegistrationHandler,
-    private val coroutineScope: CoroutineScope,
 ) : RegisterUnifiedPushUseCase {
     override 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)
         // Wait for VectorUnifiedPushMessagingReceiver.onNewEndpoint to proceed
-        return withTimeout(30.seconds) {
-            completable.await()
-        }
-            .onFailure {
-                job.cancel()
+        return runCatching {
+            withTimeout(30.seconds) {
+                val result = endpointRegistrationHandler.state
+                    .filter { it.clientSecret == clientSecret }
+                    .first()
+                    .result
+                result.getOrThrow()
             }
+        }
     }
 }
diff --git a/libraries/pushproviders/unifiedpush/src/test/kotlin/io/element/android/libraries/pushproviders/unifiedpush/DefaultRegisterUnifiedPushUseCaseTest.kt b/libraries/pushproviders/unifiedpush/src/test/kotlin/io/element/android/libraries/pushproviders/unifiedpush/DefaultRegisterUnifiedPushUseCaseTest.kt
index 99886be2dd..f898a4f804 100644
--- a/libraries/pushproviders/unifiedpush/src/test/kotlin/io/element/android/libraries/pushproviders/unifiedpush/DefaultRegisterUnifiedPushUseCaseTest.kt
+++ b/libraries/pushproviders/unifiedpush/src/test/kotlin/io/element/android/libraries/pushproviders/unifiedpush/DefaultRegisterUnifiedPushUseCaseTest.kt
@@ -25,12 +25,14 @@ import io.element.android.libraries.pushproviders.unifiedpush.registration.Endpo
 import io.element.android.libraries.pushproviders.unifiedpush.registration.RegistrationResult
 import kotlinx.coroutines.delay
 import kotlinx.coroutines.launch
+import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.advanceTimeBy
 import kotlinx.coroutines.test.runTest
-import org.junit.Ignore
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.robolectric.RobolectricTestRunner
+import kotlin.time.Duration.Companion.seconds
 
 @RunWith(RobolectricTestRunner::class)
 class DefaultRegisterUnifiedPushUseCaseTest {
@@ -64,15 +66,15 @@ class DefaultRegisterUnifiedPushUseCaseTest {
         assertThat(result.isSuccess).isFalse()
     }
 
-    @Ignore("Find a solution to test timeout")
     @Test
-    fun `test registration timeout`() = runTest {
+    fun `test registration timeout`() = runTest(StandardTestDispatcher()) {
         val endpointRegistrationHandler = EndpointRegistrationHandler()
         val useCase = createDefaultRegisterUnifiedPushUseCase(
             endpointRegistrationHandler = endpointRegistrationHandler
         )
         val aDistributor = Distributor("aValue", "aName")
         val result = useCase.execute(aDistributor, A_SECRET)
+        advanceTimeBy(30.seconds)
         assertThat(result.isSuccess).isFalse()
     }
 
@@ -83,7 +85,6 @@ class DefaultRegisterUnifiedPushUseCaseTest {
         return DefaultRegisterUnifiedPushUseCase(
             context = context,
             endpointRegistrationHandler = endpointRegistrationHandler,
-            coroutineScope = this@createDefaultRegisterUnifiedPushUseCase
         )
     }
 }