Refactor unknown state debounce

2 files changed, 28 insertions, 14 deletions

diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/MainActivity.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/MainActivity.kt
index f1831f0a88..ec31a4c706 100644
--- a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/MainActivity.kt
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/ui/MainActivity.kt
@@ -27,6 +27,8 @@ import net.mullvad.mullvadvpn.ui.fragments.DeviceRevokedFragment
 import net.mullvad.mullvadvpn.ui.serviceconnection.AccountRepository
 import net.mullvad.mullvadvpn.ui.serviceconnection.DeviceRepository
 import net.mullvad.mullvadvpn.ui.serviceconnection.ServiceConnectionManager
+import net.mullvad.mullvadvpn.util.UNKNOWN_STATE_DEBOUNCE_DELAY_MILLISECONDS
+import net.mullvad.mullvadvpn.util.addDebounceForUnknownState
 import org.koin.android.ext.android.getKoin
 import org.koin.core.context.loadKoinModules
 
@@ -159,7 +161,7 @@ open class MainActivity : FragmentActivity() {
                 .flowWithLifecycle(lifecycle, Lifecycle.State.RESUMED)
                 .debounce {
                     // Debounce DeviceState.Unknown to delay view transitions during reconnect.
-                    it.addDebounceForUnknownState()
+                    it.addDebounceForUnknownState(UNKNOWN_STATE_DEBOUNCE_DELAY_MILLISECONDS)
                 }
                 .collect { newState ->
                     if (newState != currentState) {
@@ -178,14 +180,6 @@ open class MainActivity : FragmentActivity() {
         }
     }
 
-    private fun DeviceState.addDebounceForUnknownState(): Long {
-        return if (this is DeviceState.Unknown) {
-            UNKNOWN_STATE_DEBOUNCE_DELAY_MILLISECONDS
-        } else {
-            ZERO_DEBOUNCE_DELAY_MILLISECONDS
-        }
-    }
-
     @Suppress("DEPRECATION")
     private fun requestVpnPermission() {
         val intent = VpnService.prepare(this)
@@ -243,9 +237,4 @@ open class MainActivity : FragmentActivity() {
             }
         }
     }
-
-    companion object {
-        private const val ZERO_DEBOUNCE_DELAY_MILLISECONDS = 0L
-        private const val UNKNOWN_STATE_DEBOUNCE_DELAY_MILLISECONDS = 2000L
-    }
 }
diff --git a/android/app/src/main/kotlin/net/mullvad/mullvadvpn/util/DeviceStateExtensions.kt b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/util/DeviceStateExtensions.kt
new file mode 100644
index 0000000000..d41a628bec
--- /dev/null
+++ b/android/app/src/main/kotlin/net/mullvad/mullvadvpn/util/DeviceStateExtensions.kt
@@ -0,0 +1,25 @@
+package net.mullvad.mullvadvpn.util
+
+import kotlin.reflect.KClass
+import net.mullvad.mullvadvpn.model.DeviceState
+
+const val UNKNOWN_STATE_DEBOUNCE_DELAY_MILLISECONDS = 2000L
+private const val ZERO_DEBOUNCE_DELAY_MILLISECONDS = 0L
+
+fun DeviceState.addDebounceForUnknownState(
+    delay: Long
+): Long {
+    return addDebounceForStates(delay, DeviceState.Unknown::class)
+}
+
+fun <T> DeviceState.addDebounceForStates(
+    delay: Long,
+    vararg states: KClass<T>
+): Long where T : DeviceState {
+    val result = states.any { this::class == it }
+    return if (result) {
+        delay
+    } else {
+        ZERO_DEBOUNCE_DELAY_MILLISECONDS
+    }
+}