warm up jwks cache when user enters activation flow

Author: lmac012

subscriptions/subscriptions-impl/src/main/java/com/duckduckgo/subscriptions/impl/ui/RestoreSubscriptionViewModel.kt

@@ -19,13 +19,15 @@ package com.duckduckgo.subscriptions.impl.ui
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
 import com.duckduckgo.anvil.annotations.ContributesViewModel
+import com.duckduckgo.app.di.AppCoroutineScope
 import com.duckduckgo.common.utils.DispatcherProvider
 import com.duckduckgo.di.scopes.ActivityScope
 import com.duckduckgo.subscriptions.api.SubscriptionStatus
 import com.duckduckgo.subscriptions.impl.RealSubscriptionsManager.Companion.SUBSCRIPTION_NOT_FOUND_ERROR
 import com.duckduckgo.subscriptions.impl.RealSubscriptionsManager.RecoverSubscriptionResult
 import com.duckduckgo.subscriptions.impl.SubscriptionsChecker
 import com.duckduckgo.subscriptions.impl.SubscriptionsManager
+import com.duckduckgo.subscriptions.impl.auth2.AuthClient
 import com.duckduckgo.subscriptions.impl.pixels.SubscriptionPixelSender
 import com.duckduckgo.subscriptions.impl.repository.isExpired
 import com.duckduckgo.subscriptions.impl.ui.RestoreSubscriptionViewModel.Command.Error
@@ -35,6 +37,7 @@ import com.duckduckgo.subscriptions.impl.ui.RestoreSubscriptionViewModel.Command
 import com.duckduckgo.subscriptions.impl.ui.RestoreSubscriptionViewModel.Command.SubscriptionNotFound
 import com.duckduckgo.subscriptions.impl.ui.RestoreSubscriptionViewModel.Command.Success
 import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.channels.BufferOverflow.DROP_OLDEST
 import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.flow.Flow
@@ -44,13 +47,16 @@ import kotlinx.coroutines.flow.launchIn
 import kotlinx.coroutines.flow.onEach
 import kotlinx.coroutines.flow.receiveAsFlow
 import kotlinx.coroutines.launch
+import logcat.logcat
 
 @ContributesViewModel(ActivityScope::class)
 class RestoreSubscriptionViewModel @Inject constructor(
     private val subscriptionsManager: SubscriptionsManager,
     private val subscriptionsChecker: SubscriptionsChecker,
     private val dispatcherProvider: DispatcherProvider,
     private val pixelSender: SubscriptionPixelSender,
+    private val authClient: AuthClient,
+    @AppCoroutineScope private val appCoroutineScope: CoroutineScope,
 ) : ViewModel() {
 
     private val command = Channel<Command>(1, DROP_OLDEST)
@@ -106,6 +112,7 @@ class RestoreSubscriptionViewModel @Inject constructor(
         viewModelScope.launch {
             command.send(RestoreFromEmail)
         }
+        warmUpJwksCache()
     }
 
     fun onSubscriptionRestoredFromEmail() = viewModelScope.launch {
@@ -116,6 +123,20 @@ class RestoreSubscriptionViewModel @Inject constructor(
         }
     }
 
+    /*
+        We'll need JWKs to validate auth tokens returned by FE after the user completes activation flow using email.
+        Prefetching them is optional, but it reduces the risk of failure when the network connection is unstable.
+     */
+    private fun warmUpJwksCache() {
+        appCoroutineScope.launch {
+            try {
+                authClient.getJwks()
+            } catch (e: Exception) {
+                logcat { "Failed to warm-up JWKs cache, e: ${e.stackTraceToString()}" }
+            }
+        }
+    }
+
     sealed class Command {
         data object RestoreFromEmail : Command()
         data object Success : Command()

subscriptions/subscriptions-impl/src/test/java/com/duckduckgo/subscriptions/impl/ui/RestoreSubscriptionViewModelTest.kt

@@ -10,6 +10,7 @@ import com.duckduckgo.subscriptions.impl.RealSubscriptionsManager.Companion.SUBS
 import com.duckduckgo.subscriptions.impl.RealSubscriptionsManager.RecoverSubscriptionResult
 import com.duckduckgo.subscriptions.impl.SubscriptionsChecker
 import com.duckduckgo.subscriptions.impl.SubscriptionsManager
+import com.duckduckgo.subscriptions.impl.auth2.AuthClient
 import com.duckduckgo.subscriptions.impl.pixels.SubscriptionPixelSender
 import com.duckduckgo.subscriptions.impl.repository.Subscription
 import com.duckduckgo.subscriptions.impl.ui.RestoreSubscriptionViewModel.Command.Error
@@ -37,6 +38,7 @@ class RestoreSubscriptionViewModelTest {
     private val subscriptionsManager: SubscriptionsManager = mock()
     private val pixelSender: SubscriptionPixelSender = mock()
     private val subscriptionsChecker: SubscriptionsChecker = mock()
+    private val authClient: AuthClient = mock()
     private lateinit var viewModel: RestoreSubscriptionViewModel
 
     @Before
@@ -46,6 +48,8 @@ class RestoreSubscriptionViewModelTest {
             dispatcherProvider = coroutineTestRule.testDispatcherProvider,
             pixelSender = pixelSender,
             subscriptionsChecker = subscriptionsChecker,
+            authClient = authClient,
+            appCoroutineScope = coroutineTestRule.testScope,
         )
     }
 
@@ -190,6 +194,24 @@ class RestoreSubscriptionViewModelTest {
         }
     }
 
+    @Test
+    fun whenRestoreFromEmailThenJwksCacheIsWarmedUp() = runTest {
+        viewModel.restoreFromEmail()
+        verify(authClient).getJwks()
+    }
+
+    @Test
+    fun whenWarmUpJwksFailsThenNoCrashOccurs() = runTest {
+        whenever(authClient.getJwks()).thenThrow(RuntimeException("Network error"))
+
+        viewModel.restoreFromEmail()
+
+        viewModel.commands().test {
+            assertTrue(awaitItem() is RestoreFromEmail)
+        }
+        verify(pixelSender).reportActivateSubscriptionEnterEmailClick()
+    }
+
     private fun subscriptionActive(): Subscription {
         return Subscription(
             productId = "productId",