Repository: element-hq/element-x-android (language: kotlin)

Issue:
Push-Notification Option for recreate token if test failed
Here a dedicated issue of this slice different Problem:
https://github.com/element-hq/element-x-android/issues/3570#issuecomment-2423760827

### Steps to reproduce

1. Login to EXA
2. See in microg if EXA was correct registerd to FCM
3. test push works in EXA
4. Delete registration in microg
5. Run test in EXA again


### Outcome

#### What did you expect?

Token validation success, but loop test failed.

#### What happened instead?
Here are two Option to fix it:
 - Option to reset the token to trigger new FCM registration
 - Reset FCM-Token after change Norification Provider (e.g. to ntfy)

### Your phone model

_No response_

### Operating system version

_No response_

### Application version and app store

_No response_

### Homeserver

_No response_

### Will you send logs?

No

### Are you willing to provide a PR?

No


Discussion hints:
Here a Screenshot of current broken state

![1000008557](https://github.com/user-attachments/assets/eb849f57-1daa-446e-9214-0fce482b24e9)


Pull request: Rotate firebase token in case of error
<!-- Please read [CONTRIBUTING.md](https://github.com/element-hq/element-x-android/blob/develop/CONTRIBUTING.md) before submitting your pull request -->
 
## Content

<!-- Describe shortly what has been changed -->
Add a quick fix for the troubleshoot notification test in case of PusherRejected error, which means that the token is unknown.

In this case, the quick fix will rotate the Firebase token.

Note: maybe later we will see if we can do the same thing for UnifiedPush provider.

## Motivation and context

<!-- Provide link to the corresponding issue if applicable or explain the context -->
Closes #3752 

## Screenshots / GIFs

<!--
We have screenshot tests in the project, so attaching screenshots to a PR is not mandatory, as far as there
is a Composable Preview covering the changes. In this case, the change will appear in the file diff.
Note that all the UI composables should be covered by a Composable Preview.

Providing a video of the change is still very useful for the reviewer and for the history of the project.

You can use a table like this to show screenshots comparison.
Uncomment this markdown table below and edit the last line `|||`:
|copy screenshot of before here|copy screenshot of after here|

|Before|After|
|-|-|
|||
 -->

You can see in the video below that when clicking on QuickFix, the token is rotated.

https://github.com/user-attachments/assets/0152344f-809e-482b-9cfa-1e2a6673d54d

## Tests

<!-- Explain how you tested your development -->

- @wrenix can you check if this is fixing your issue please?

Internally, this can be tested by patching the code and let `DefaultPushService.testPush()` throw a `PushGatewayFailure.PusherRejected()` (this is what I have done for the video above).

## Tested devices

- [ ] Physical
- [x] Emulator
- OS version(s):

## Checklist

<!-- Depending on the Pull Request content, it can be acceptable if some of the following checkboxes stay unchecked. -->

- [ ] Changes have been tested on an Android device or Android emulator with API 23
- [ ] UI change has been tested on both light and dark themes
- [ ] Accessibility has been taken into account. See https://github.com/element-hq/element-x-android/blob/develop/CONTRIBUTING.md#accessibility
- [ ] Pull request is based on the develop branch
- [ ] Pull request title will be used in the release note, it clearly define what will change for the user
- [ ] Pull request includes screenshots or videos if containing UI changes
- [ ] Pull request includes a [sign off](https://matrix-org.github.io/synapse/latest/development/contributing_guide.html#sign-off)
- [ ] You've made a self review of your PR


Test command: ./gradlew :libraries:pushproviders:firebase:testDebugUnitTest --tests "io.element.android.libraries.pushproviders.firebase.troubleshoot.FirebaseTokenTestTest" --tests "io.element.android.libraries.pushproviders.firebase.FirebasePushProviderTest" :libraries:push:impl:testDebugUnitTest --tests "io.element.android.libraries.push.impl.troubleshoot.PushLoopbackTestTest"

Fix patch:
diff --git a/libraries/pushproviders/api/src/main/kotlin/io/element/android/libraries/pushproviders/api/PushProvider.kt b/libraries/pushproviders/api/src/main/kotlin/io/element/android/libraries/pushproviders/api/PushProvider.kt
index 0f3cde8033b..90b4cb04658 100644
--- a/libraries/pushproviders/api/src/main/kotlin/io/element/android/libraries/pushproviders/api/PushProvider.kt
+++ b/libraries/pushproviders/api/src/main/kotlin/io/element/android/libraries/pushproviders/api/PushProvider.kt
@@ -44,4 +44,10 @@ interface PushProvider {
     suspend fun unregister(matrixClient: MatrixClient): Result<Unit>
 
     suspend fun getCurrentUserPushConfig(): CurrentUserPushConfig?
+
+    fun canRotateToken(): Boolean
+
+    suspend fun rotateToken(): Result<Unit> {
+        error("rotateToken() not implemented, you need to override this method in your implementation")
+    }
 }
diff --git a/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/FirebasePushProvider.kt b/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/FirebasePushProvider.kt
index 6bbdb1efe83..b1e31a23035 100644
--- a/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/FirebasePushProvider.kt
+++ b/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/FirebasePushProvider.kt
@@ -25,6 +25,7 @@ class FirebasePushProvider @Inject constructor(
     private val firebaseStore: FirebaseStore,
     private val pusherSubscriber: PusherSubscriber,
     private val isPlayServiceAvailable: IsPlayServiceAvailable,
+    private val firebaseTokenRotator: FirebaseTokenRotator,
 ) : PushProvider {
     override val index = FirebaseConfig.INDEX
     override val name = FirebaseConfig.NAME
@@ -71,6 +72,12 @@ class FirebasePushProvider @Inject constructor(
         }
     }
 
+    override fun canRotateToken(): Boolean = true
+
+    override suspend fun rotateToken(): Result<Unit> {
+        return firebaseTokenRotator.rotate()
+    }
+
     companion object {
         private val firebaseDistributor = Distributor("Firebase", "Firebase")
     }
diff --git a/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/FirebaseStore.kt b/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/FirebaseStore.kt
index 2058f716d4e..4ad36194544 100644
--- a/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/FirebaseStore.kt
+++ b/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/FirebaseStore.kt
@@ -11,6 +11,10 @@ import android.content.SharedPreferences
 import androidx.core.content.edit
 import com.squareup.anvil.annotations.ContributesBinding
 import io.element.android.libraries.di.AppScope
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
+import kotlinx.coroutines.flow.onCompletion
+import kotlinx.coroutines.flow.onStart
 import javax.inject.Inject
 
 /**
@@ -18,6 +22,7 @@ import javax.inject.Inject
  */
 interface FirebaseStore {
     fun getFcmToken(): String?
+    fun fcmTokenFlow(): Flow<String?>
     fun storeFcmToken(token: String?)
 }
 
@@ -29,6 +34,22 @@ class SharedPreferencesFirebaseStore @Inject constructor(
         return sharedPreferences.getString(PREFS_KEY_FCM_TOKEN, null)
     }
 
+    override fun fcmTokenFlow(): Flow<String?> {
+        val flow = MutableStateFlow(getFcmToken())
+        val listener = SharedPreferences.OnSharedPreferenceChangeListener { _, k ->
+            if (k == PREFS_KEY_FCM_TOKEN) {
+                try {
+                    flow.value = getFcmToken()
+                } catch (e: Exception) {
+                    flow.value = null
+                }
+            }
+        }
+        return flow
+            .onStart { sharedPreferences.registerOnSharedPreferenceChangeListener(listener) }
+            .onCompletion { sharedPreferences.unregisterOnSharedPreferenceChangeListener(listener) }
+    }
+
     override fun storeFcmToken(token: String?) {
         sharedPreferences.edit {
             putString(PREFS_KEY_FCM_TOKEN, token)
diff --git a/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/FirebaseTokenDeleter.kt b/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/FirebaseTokenDeleter.kt
new file mode 100644
index 00000000000..597a39493f8
--- /dev/null
+++ b/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/FirebaseTokenDeleter.kt
@@ -0,0 +1,49 @@
+/*
+ * Copyright 2023, 2024 New Vector Ltd.
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only
+ * Please see LICENSE in the repository root for full details.
+ */
+
+package io.element.android.libraries.pushproviders.firebase
+
+import com.google.firebase.messaging.FirebaseMessaging
+import com.squareup.anvil.annotations.ContributesBinding
+import io.element.android.libraries.di.AppScope
+import timber.log.Timber
+import javax.inject.Inject
+import kotlin.coroutines.resume
+import kotlin.coroutines.resumeWithException
+import kotlin.coroutines.suspendCoroutine
+
+interface FirebaseTokenDeleter {
+    /**
+     * Deletes the current Firebase token.
+     */
+    suspend fun delete()
+}
+
+@ContributesBinding(AppScope::class)
+class DefaultFirebaseTokenDeleter @Inject constructor(
+    private val isPlayServiceAvailable: IsPlayServiceAvailable,
+) : FirebaseTokenDeleter {
+    override suspend fun delete() {
+        // 'app should always check the device for a compatible Google Play services APK before accessing Google Play services features'
+        isPlayServiceAvailable.checkAvailableOrThrow()
+        suspendCoroutine { continuation ->
+            try {
+                FirebaseMessaging.getInstance().deleteToken()
+                    .addOnSuccessListener {
+                        continuation.resume(Unit)
+                    }
+                    .addOnFailureListener { e ->
+                        Timber.e(e, "## deleteFirebaseToken() : failed")
+                        continuation.resumeWithException(e)
+                    }
+            } catch (e: Throwable) {
+                Timber.e(e, "## deleteFirebaseToken() : failed")
+                continuation.resumeWithException(e)
+            }
+        }
+    }
+}
diff --git a/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/FirebaseTokenGetter.kt b/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/FirebaseTokenGetter.kt
new file mode 100644
index 00000000000..a15d5a250c1
--- /dev/null
+++ b/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/FirebaseTokenGetter.kt
@@ -0,0 +1,50 @@
+/*
+ * Copyright 2023, 2024 New Vector Ltd.
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only
+ * Please see LICENSE in the repository root for full details.
+ */
+
+package io.element.android.libraries.pushproviders.firebase
+
+import com.google.firebase.messaging.FirebaseMessaging
+import com.squareup.anvil.annotations.ContributesBinding
+import io.element.android.libraries.di.AppScope
+import timber.log.Timber
+import javax.inject.Inject
+import kotlin.coroutines.resume
+import kotlin.coroutines.resumeWithException
+import kotlin.coroutines.suspendCoroutine
+
+interface FirebaseTokenGetter {
+    /**
+     * Read the current Firebase token from FirebaseMessaging.
+     * If the token does not exist, it will be generated.
+     */
+    suspend fun get(): String
+}
+
+@ContributesBinding(AppScope::class)
+class DefaultFirebaseTokenGetter @Inject constructor(
+    private val isPlayServiceAvailable: IsPlayServiceAvailable,
+) : FirebaseTokenGetter {
+    override suspend fun get(): String {
+        // 'app should always check the device for a compatible Google Play services APK before accessing Google Play services features'
+        isPlayServiceAvailable.checkAvailableOrThrow()
+        return suspendCoroutine { continuation ->
+            try {
+                FirebaseMessaging.getInstance().token
+                    .addOnSuccessListener { token ->
+                        continuation.resume(token)
+                    }
+                    .addOnFailureListener { e ->
+                        Timber.e(e, "## retrievedFirebaseToken() : failed")
+                        continuation.resumeWithException(e)
+                    }
+            } catch (e: Throwable) {
+                Timber.e(e, "## retrievedFirebaseToken() : failed")
+                continuation.resumeWithException(e)
+            }
+        }
+    }
+}
diff --git a/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/FirebaseTokenRotator.kt b/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/FirebaseTokenRotator.kt
new file mode 100644
index 00000000000..fc55d11a77a
--- /dev/null
+++ b/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/FirebaseTokenRotator.kt
@@ -0,0 +1,32 @@
+/*
+ * Copyright 2024 New Vector Ltd.
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only
+ * Please see LICENSE in the repository root for full details.
+ */
+
+package io.element.android.libraries.pushproviders.firebase
+
+import com.squareup.anvil.annotations.ContributesBinding
+import io.element.android.libraries.di.AppScope
+import javax.inject.Inject
+
+interface FirebaseTokenRotator {
+    suspend fun rotate(): Result<Unit>
+}
+
+/**
+ * This class delete the Firebase token and generate a new one.
+ */
+@ContributesBinding(AppScope::class)
+class DefaultFirebaseTokenRotator @Inject constructor(
+    private val firebaseTokenDeleter: FirebaseTokenDeleter,
+    private val firebaseTokenGetter: FirebaseTokenGetter,
+) : FirebaseTokenRotator {
+    override suspend fun rotate(): Result<Unit> {
+        return runCatching {
+            firebaseTokenDeleter.delete()
+            firebaseTokenGetter.get()
+        }
+    }
+}
diff --git a/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/FirebaseTroubleshooter.kt b/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/FirebaseTroubleshooter.kt
index 6ef6931ea14..6991da5f251 100644
--- a/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/FirebaseTroubleshooter.kt
+++ b/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/FirebaseTroubleshooter.kt
@@ -7,14 +7,9 @@
 
 package io.element.android.libraries.pushproviders.firebase
 
-import com.google.firebase.messaging.FirebaseMessaging
 import com.squareup.anvil.annotations.ContributesBinding
 import io.element.android.libraries.di.AppScope
-import timber.log.Timber
 import javax.inject.Inject
-import kotlin.coroutines.resume
-import kotlin.coroutines.resumeWithException
-import kotlin.coroutines.suspendCoroutine
 
 interface FirebaseTroubleshooter {
     suspend fun troubleshoot(): Result<Unit>
@@ -26,37 +21,12 @@ interface FirebaseTroubleshooter {
 @ContributesBinding(AppScope::class)
 class DefaultFirebaseTroubleshooter @Inject constructor(
     private val newTokenHandler: FirebaseNewTokenHandler,
-    private val isPlayServiceAvailable: IsPlayServiceAvailable,
+    private val firebaseTokenGetter: FirebaseTokenGetter,
 ) : FirebaseTroubleshooter {
     override suspend fun troubleshoot(): Result<Unit> {
         return runCatching {
-            val token = retrievedFirebaseToken()
+            val token = firebaseTokenGetter.get()
             newTokenHandler.handle(token)
         }
     }
-
-    private suspend fun retrievedFirebaseToken(): String {
-        return suspendCoroutine { continuation ->
-            // 'app should always check the device for a compatible Google Play services APK before accessing Google Play services features'
-            if (isPlayServiceAvailable.isAvailable()) {
-                try {
-                    FirebaseMessaging.getInstance().token
-                        .addOnSuccessListener { token ->
-                            continuation.resume(token)
-                        }
-                        .addOnFailureListener { e ->
-                            Timber.e(e, "## retrievedFirebaseToken() : failed")
-                            continuation.resumeWithException(e)
-                        }
-                } catch (e: Throwable) {
-                    Timber.e(e, "## retrievedFirebaseToken() : failed")
-                    continuation.resumeWithException(e)
-                }
-            } else {
-                val e = Exception("No valid Google Play Services found. Cannot use FCM.")
-                Timber.e(e)
-                continuation.resumeWithException(e)
-            }
-        }
-    }
 }
diff --git a/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/IsPlayServiceAvailable.kt b/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/IsPlayServiceAvailable.kt
index 02cf66cb744..a0177ef7576 100644
--- a/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/IsPlayServiceAvailable.kt
+++ b/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/IsPlayServiceAvailable.kt
@@ -20,6 +20,12 @@ interface IsPlayServiceAvailable {
     fun isAvailable(): Boolean
 }
 
+fun IsPlayServiceAvailable.checkAvailableOrThrow() {
+    if (!isAvailable()) {
+        throw Exception("No valid Google Play Services found. Cannot use FCM.").also(Timber::e)
+    }
+}
+
 @ContributesBinding(AppScope::class)
 class DefaultIsPlayServiceAvailable @Inject constructor(
     @ApplicationContext private val context: Context,
diff --git a/libraries/pushproviders/unifiedpush/src/main/kotlin/io/element/android/libraries/pushproviders/unifiedpush/UnifiedPushProvider.kt b/libraries/pushproviders/unifiedpush/src/main/kotlin/io/element/android/libraries/pushproviders/unifiedpush/UnifiedPushProvider.kt
index 9a2d5d5c10f..1ce929a75b6 100644
--- a/libraries/pushproviders/unifiedpush/src/main/kotlin/io/element/android/libraries/pushproviders/unifiedpush/UnifiedPushProvider.kt
+++ b/libraries/pushproviders/unifiedpush/src/main/kotlin/io/element/android/libraries/pushproviders/unifiedpush/UnifiedPushProvider.kt
@@ -53,4 +53,6 @@ class UnifiedPushProvider @Inject constructor(
     override suspend fun getCurrentUserPushConfig(): CurrentUserPushConfig? {
         return unifiedPushCurrentUserPushConfigProvider.provide()
     }
+
+    override fun canRotateToken(): Boolean = false
 }


Test patch:
diff --git a/libraries/push/impl/src/main/kotlin/io/element/android/libraries/push/impl/troubleshoot/PushLoopbackTest.kt b/libraries/push/impl/src/main/kotlin/io/element/android/libraries/push/impl/troubleshoot/PushLoopbackTest.kt
index e91f4351f78..57c255a4064 100644
--- a/libraries/push/impl/src/main/kotlin/io/element/android/libraries/push/impl/troubleshoot/PushLoopbackTest.kt
+++ b/libraries/push/impl/src/main/kotlin/io/element/android/libraries/push/impl/troubleshoot/PushLoopbackTest.kt
@@ -52,9 +52,10 @@ class PushLoopbackTest @Inject constructor(
         val testPushResult = try {
             pushService.testPush()
         } catch (pusherRejected: PushGatewayFailure.PusherRejected) {
+            val hasQuickFix = pushService.getCurrentPushProvider()?.canRotateToken() == true
             delegate.updateState(
                 description = stringProvider.getString(R.string.troubleshoot_notifications_test_push_loop_back_failure_1),
-                status = NotificationTroubleshootTestState.Status.Failure(false)
+                status = NotificationTroubleshootTestState.Status.Failure(hasQuickFix)
             )
             job.cancel()
             return
@@ -96,5 +97,11 @@ class PushLoopbackTest @Inject constructor(
         )
     }
 
+    override suspend fun quickFix(coroutineScope: CoroutineScope) {
+        delegate.start()
+        pushService.getCurrentPushProvider()?.rotateToken()
+        run(coroutineScope)
+    }
+
     override suspend fun reset() = delegate.reset()
 }
diff --git a/libraries/push/impl/src/test/kotlin/io/element/android/libraries/push/impl/troubleshoot/PushLoopbackTestTest.kt b/libraries/push/impl/src/test/kotlin/io/element/android/libraries/push/impl/troubleshoot/PushLoopbackTestTest.kt
index b12e0cf80b7..57ba07a5f64 100644
--- a/libraries/push/impl/src/test/kotlin/io/element/android/libraries/push/impl/troubleshoot/PushLoopbackTestTest.kt
+++ b/libraries/push/impl/src/test/kotlin/io/element/android/libraries/push/impl/troubleshoot/PushLoopbackTestTest.kt
@@ -13,9 +13,11 @@ import io.element.android.libraries.matrix.test.AN_EXCEPTION
 import io.element.android.libraries.matrix.test.A_FAILURE_REASON
 import io.element.android.libraries.push.api.gateway.PushGatewayFailure
 import io.element.android.libraries.push.test.FakePushService
+import io.element.android.libraries.pushproviders.test.FakePushProvider
 import io.element.android.libraries.troubleshoot.api.test.NotificationTroubleshootTestState
 import io.element.android.services.toolbox.test.strings.FakeStringProvider
 import io.element.android.services.toolbox.test.systemclock.FakeSystemClock
+import io.element.android.tests.testutils.lambda.lambdaRecorder
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.test.runTest
 import org.junit.Test
@@ -67,6 +69,41 @@ class PushLoopbackTestTest {
         }
     }
 
+    @Test
+    fun `test PushLoopbackTest PusherRejected error with quick fix`() = runTest {
+        val diagnosticPushHandler = DiagnosticPushHandler()
+        val rotateTokenLambda = lambdaRecorder<Result<Unit>> { Result.success(Unit) }
+        val sut = PushLoopbackTest(
+            pushService = FakePushService(
+                testPushBlock = {
+                    throw PushGatewayFailure.PusherRejected()
+                },
+                currentPushProvider = {
+                    FakePushProvider(
+                        canRotateTokenResult = { true },
+                        rotateTokenLambda = rotateTokenLambda,
+                    )
+                }
+            ),
+            diagnosticPushHandler = diagnosticPushHandler,
+            clock = FakeSystemClock(),
+            stringProvider = FakeStringProvider(),
+        )
+        launch {
+            sut.run(this)
+        }
+        sut.state.test {
+            assertThat(awaitItem().status).isEqualTo(NotificationTroubleshootTestState.Status.Idle(true))
+            assertThat(awaitItem().status).isEqualTo(NotificationTroubleshootTestState.Status.InProgress)
+            val lastItem = awaitItem()
+            assertThat(lastItem.status).isEqualTo(NotificationTroubleshootTestState.Status.Failure(true))
+            sut.quickFix(this)
+            assertThat(awaitItem().status).isEqualTo(NotificationTroubleshootTestState.Status.InProgress)
+            assertThat(awaitItem().status).isEqualTo(NotificationTroubleshootTestState.Status.Failure(true))
+            rotateTokenLambda.assertions().isCalledOnce()
+        }
+    }
+
     @Test
     fun `test PushLoopbackTest setup error`() = runTest {
         val diagnosticPushHandler = DiagnosticPushHandler()
diff --git a/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/troubleshoot/FirebaseTokenTest.kt b/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/troubleshoot/FirebaseTokenTest.kt
index a0d6d22d101..7d042396697 100644
--- a/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/troubleshoot/FirebaseTokenTest.kt
+++ b/libraries/pushproviders/firebase/src/main/kotlin/io/element/android/libraries/pushproviders/firebase/troubleshoot/FirebaseTokenTest.kt
@@ -19,7 +19,10 @@ import io.element.android.libraries.troubleshoot.api.test.NotificationTroublesho
 import io.element.android.libraries.troubleshoot.api.test.TestFilterData
 import io.element.android.services.toolbox.api.strings.StringProvider
 import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Job
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.launchIn
+import kotlinx.coroutines.flow.onEach
 import javax.inject.Inject
 
 @ContributesMultibinding(AppScope::class)
@@ -41,23 +44,28 @@ class FirebaseTokenTest @Inject constructor(
         return data.currentPushProviderName == FirebaseConfig.NAME
     }
 
+    private var currentJob: Job? = null
     override suspend fun run(coroutineScope: CoroutineScope) {
+        currentJob?.cancel()
         delegate.start()
-        val token = firebaseStore.getFcmToken()
-        if (token != null) {
-            delegate.updateState(
-                description = stringProvider.getString(
-                    R.string.troubleshoot_notifications_test_firebase_token_success,
-                    "${token.take(8)}*****"
-                ),
-                status = NotificationTroubleshootTestState.Status.Success
-            )
-        } else {
-            delegate.updateState(
-                description = stringProvider.getString(R.string.troubleshoot_notifications_test_firebase_token_failure),
-                status = NotificationTroubleshootTestState.Status.Failure(true)
-            )
-        }
+        currentJob = firebaseStore.fcmTokenFlow()
+            .onEach { token ->
+                if (token != null) {
+                    delegate.updateState(
+                        description = stringProvider.getString(
+                            R.string.troubleshoot_notifications_test_firebase_token_success,
+                            "*****${token.takeLast(8)}"
+                        ),
+                        status = NotificationTroubleshootTestState.Status.Success
+                    )
+                } else {
+                    delegate.updateState(
+                        description = stringProvider.getString(R.string.troubleshoot_notifications_test_firebase_token_failure),
+                        status = NotificationTroubleshootTestState.Status.Failure(true)
+                    )
+                }
+            }
+            .launchIn(coroutineScope)
     }
 
     override suspend fun reset() = delegate.reset()
diff --git a/libraries/pushproviders/firebase/src/test/kotlin/io/element/android/libraries/pushproviders/firebase/FakeFirebaseTokenRotator.kt b/libraries/pushproviders/firebase/src/test/kotlin/io/element/android/libraries/pushproviders/firebase/FakeFirebaseTokenRotator.kt
new file mode 100644
index 00000000000..c96c67c3fa4
--- /dev/null
+++ b/libraries/pushproviders/firebase/src/test/kotlin/io/element/android/libraries/pushproviders/firebase/FakeFirebaseTokenRotator.kt
@@ -0,0 +1,18 @@
+/*
+ * Copyright 2024 New Vector Ltd.
+ *
+ * SPDX-License-Identifier: AGPL-3.0-only
+ * Please see LICENSE in the repository root for full details.
+ */
+
+package io.element.android.libraries.pushproviders.firebase
+
+import io.element.android.tests.testutils.lambda.lambdaError
+
+class FakeFirebaseTokenRotator(
+    private val rotateWithResult: () -> Result<Unit> = { lambdaError() }
+) : FirebaseTokenRotator {
+    override suspend fun rotate(): Result<Unit> {
+        return rotateWithResult()
+    }
+}
diff --git a/libraries/pushproviders/firebase/src/test/kotlin/io/element/android/libraries/pushproviders/firebase/FirebasePushProviderTest.kt b/libraries/pushproviders/firebase/src/test/kotlin/io/element/android/libraries/pushproviders/firebase/FirebasePushProviderTest.kt
index f94f4ebac89..6a3a9474e40 100644
--- a/libraries/pushproviders/firebase/src/test/kotlin/io/element/android/libraries/pushproviders/firebase/FirebasePushProviderTest.kt
+++ b/libraries/pushproviders/firebase/src/test/kotlin/io/element/android/libraries/pushproviders/firebase/FirebasePushProviderTest.kt
@@ -166,15 +166,27 @@ class FirebasePushProviderTest {
         assertThat(result).isEqualTo(CurrentUserPushConfig(FirebaseConfig.PUSHER_HTTP_URL, "aToken"))
     }
 
+    @Test
+    fun `rotateToken invokes the FirebaseTokenRotator`() = runTest {
+        val lambda = lambdaRecorder<Result<Unit>> { Result.success(Unit) }
+        val firebasePushProvider = createFirebasePushProvider(
+            firebaseTokenRotator = FakeFirebaseTokenRotator(lambda),
+        )
+        firebasePushProvider.rotateToken()
+        lambda.assertions().isCalledOnce()
+    }
+
     private fun createFirebasePushProvider(
         firebaseStore: FirebaseStore = InMemoryFirebaseStore(),
         pusherSubscriber: PusherSubscriber = FakePusherSubscriber(),
         isPlayServiceAvailable: IsPlayServiceAvailable = FakeIsPlayServiceAvailable(false),
+        firebaseTokenRotator: FirebaseTokenRotator = FakeFirebaseTokenRotator(),
     ): FirebasePushProvider {
         return FirebasePushProvider(
             firebaseStore = firebaseStore,
             pusherSubscriber = pusherSubscriber,
             isPlayServiceAvailable = isPlayServiceAvailable,
+            firebaseTokenRotator = firebaseTokenRotator,
         )
     }
 }
diff --git a/libraries/pushproviders/firebase/src/test/kotlin/io/element/android/libraries/pushproviders/firebase/InMemoryFirebaseStore.kt b/libraries/pushproviders/firebase/src/test/kotlin/io/element/android/libraries/pushproviders/firebase/InMemoryFirebaseStore.kt
index 3a3e582f7f2..6728fd04903 100644
--- a/libraries/pushproviders/firebase/src/test/kotlin/io/element/android/libraries/pushproviders/firebase/InMemoryFirebaseStore.kt
+++ b/libraries/pushproviders/firebase/src/test/kotlin/io/element/android/libraries/pushproviders/firebase/InMemoryFirebaseStore.kt
@@ -7,11 +7,16 @@
 
 package io.element.android.libraries.pushproviders.firebase
 
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.flowOf
+
 class InMemoryFirebaseStore(
     private var token: String? = null
 ) : FirebaseStore {
     override fun getFcmToken(): String? = token
 
+    override fun fcmTokenFlow(): Flow<String?> = flowOf(token)
+
     override fun storeFcmToken(token: String?) {
         this.token = token
     }
diff --git a/libraries/pushproviders/firebase/src/test/kotlin/io/element/android/libraries/pushproviders/firebase/troubleshoot/FirebaseTokenTestTest.kt b/libraries/pushproviders/firebase/src/test/kotlin/io/element/android/libraries/pushproviders/firebase/troubleshoot/FirebaseTokenTestTest.kt
index adf50300e77..e67c59c8867 100644
--- a/libraries/pushproviders/firebase/src/test/kotlin/io/element/android/libraries/pushproviders/firebase/troubleshoot/FirebaseTokenTestTest.kt
+++ b/libraries/pushproviders/firebase/src/test/kotlin/io/element/android/libraries/pushproviders/firebase/troubleshoot/FirebaseTokenTestTest.kt
@@ -35,7 +35,7 @@ class FirebaseTokenTestTest {
             assertThat(awaitItem().status).isEqualTo(NotificationTroubleshootTestState.Status.InProgress)
             val lastItem = awaitItem()
             assertThat(lastItem.status).isEqualTo(NotificationTroubleshootTestState.Status.Success)
-            assertThat(lastItem.description).contains(FAKE_TOKEN.take(8))
+            assertThat(lastItem.description).contains(FAKE_TOKEN.takeLast(8))
             assertThat(lastItem.description).doesNotContain(FAKE_TOKEN)
         }
     }
diff --git a/libraries/pushproviders/test/src/main/kotlin/io/element/android/libraries/pushproviders/test/FakePushProvider.kt b/libraries/pushproviders/test/src/main/kotlin/io/element/android/libraries/pushproviders/test/FakePushProvider.kt
index 9fe222e1cba..6f5bdda5b77 100644
--- a/libraries/pushproviders/test/src/main/kotlin/io/element/android/libraries/pushproviders/test/FakePushProvider.kt
+++ b/libraries/pushproviders/test/src/main/kotlin/io/element/android/libraries/pushproviders/test/FakePushProvider.kt
@@ -21,6 +21,8 @@ class FakePushProvider(
     private val currentUserPushConfig: CurrentUserPushConfig? = null,
     private val registerWithResult: (MatrixClient, Distributor) -> Result<Unit> = { _, _ -> lambdaError() },
     private val unregisterWithResult: (MatrixClient) -> Result<Unit> = { lambdaError() },
+    private val canRotateTokenResult: () -> Boolean = { lambdaError() },
+    private val rotateTokenLambda: () -> Result<Unit> = { lambdaError() },
 ) : PushProvider {
     override fun getDistributors(): List<Distributor> = distributors
 
@@ -39,4 +41,12 @@ class FakePushProvider(
     override suspend fun getCurrentUserPushConfig(): CurrentUserPushConfig? {
         return currentUserPushConfig
     }
+
+    override fun canRotateToken(): Boolean {
+        return canRotateTokenResult()
+    }
+
+    override suspend fun rotateToken(): Result<Unit> {
+        return rotateTokenLambda()
+    }
 }