Wallet quick affordance binder calls on main thread

Move several calls to the background thread that were causing jank

Fixes: 312450479
Flag: NONE
Test: atest QuickAccessWalletKeyguardQuickAffordanceConfigTest.kt
Test: manual - used wallet sandbox environment to validate changes
Change-Id: Ia0db621c1904ee8a59413e890c5ddd92a7a2783b
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/quickaffordance/QuickAccessWalletKeyguardQuickAffordanceConfigTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/quickaffordance/QuickAccessWalletKeyguardQuickAffordanceConfigTest.kt
index 02db0d7..a613ad8 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/quickaffordance/QuickAccessWalletKeyguardQuickAffordanceConfigTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyguard/data/quickaffordance/QuickAccessWalletKeyguardQuickAffordanceConfigTest.kt
@@ -29,17 +29,16 @@
 import com.android.systemui.animation.Expandable
 import com.android.systemui.common.shared.model.ContentDescription
 import com.android.systemui.common.shared.model.Icon
+import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.plugins.ActivityStarter
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.mock
 import com.android.systemui.util.mockito.whenever
 import com.android.systemui.wallet.controller.QuickAccessWalletController
 import com.google.common.truth.Truth.assertThat
-import kotlinx.coroutines.ExperimentalCoroutinesApi
-import kotlinx.coroutines.flow.launchIn
-import kotlinx.coroutines.flow.onEach
-import kotlinx.coroutines.test.UnconfinedTestDispatcher
-import kotlinx.coroutines.test.runBlockingTest
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.TestDispatcher
+import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.runTest
 import org.junit.Before
 import org.junit.Test
@@ -48,7 +47,6 @@
 import org.mockito.Mockito.verify
 import org.mockito.MockitoAnnotations
 
-@OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
 @RunWith(AndroidJUnit4::class)
 class QuickAccessWalletKeyguardQuickAffordanceConfigTest : SysuiTestCase() {
@@ -56,26 +54,32 @@
     @Mock private lateinit var walletController: QuickAccessWalletController
     @Mock private lateinit var activityStarter: ActivityStarter
 
+    private lateinit var testDispatcher: TestDispatcher
+    private lateinit var testScope: TestScope
+
     private lateinit var underTest: QuickAccessWalletKeyguardQuickAffordanceConfig
 
     @Before
     fun setUp() {
         MockitoAnnotations.initMocks(this)
 
+        testDispatcher = StandardTestDispatcher()
+        testScope = TestScope(testDispatcher)
+
         underTest =
             QuickAccessWalletKeyguardQuickAffordanceConfig(
                 context,
+                testDispatcher,
                 walletController,
                 activityStarter,
             )
     }
 
     @Test
-    fun affordance_keyguardShowing_hasWalletCard_visibleModel() = runBlockingTest {
+    fun affordance_keyguardShowing_hasWalletCard_visibleModel() = testScope.runTest {
         setUpState()
-        var latest: KeyguardQuickAffordanceConfig.LockScreenState? = null
 
-        val job = underTest.lockScreenState.onEach { latest = it }.launchIn(this)
+        val latest by collectLastValue(underTest.lockScreenState)
 
         val visibleModel = latest as KeyguardQuickAffordanceConfig.LockScreenState.Visible
         assertThat(visibleModel.icon)
@@ -88,77 +92,61 @@
                         ),
                 )
             )
-        job.cancel()
     }
 
     @Test
-    fun affordance_keyguardShowing_hasNonPaymentCard_modelIsNone() =
-        runTest(UnconfinedTestDispatcher()) {
+    fun affordance_keyguardShowing_hasNonPaymentCard_modelIsNone() = testScope.runTest {
             setUpState(cardType = WalletCard.CARD_TYPE_NON_PAYMENT)
-            var latest: KeyguardQuickAffordanceConfig.LockScreenState? = null
 
-            val job = underTest.lockScreenState.onEach { latest = it }.launchIn(this)
+            val latest by collectLastValue(underTest.lockScreenState)
 
             assertThat(latest).isEqualTo(KeyguardQuickAffordanceConfig.LockScreenState.Hidden)
-            job.cancel()
         }
 
     @Test
-    fun affordance_keyguardShowing_hasPaymentCard_visibleModel() =
-        runTest(UnconfinedTestDispatcher()) {
-            setUpState(cardType = WalletCard.CARD_TYPE_PAYMENT)
-            var latest: KeyguardQuickAffordanceConfig.LockScreenState? = null
+    fun affordance_keyguardShowing_hasPaymentCard_visibleModel() = testScope.runTest {
+        setUpState(cardType = WalletCard.CARD_TYPE_PAYMENT)
 
-            val job = underTest.lockScreenState.onEach { latest = it }.launchIn(this)
+        val latest by collectLastValue(underTest.lockScreenState)
 
-            val visibleModel = latest as KeyguardQuickAffordanceConfig.LockScreenState.Visible
-            assertThat(visibleModel.icon)
-                .isEqualTo(
-                    Icon.Loaded(
-                        drawable = ICON,
-                        contentDescription =
-                            ContentDescription.Resource(
-                                res = R.string.accessibility_wallet_button,
-                            ),
-                    )
+        val visibleModel = latest as KeyguardQuickAffordanceConfig.LockScreenState.Visible
+        assertThat(visibleModel.icon)
+            .isEqualTo(
+                Icon.Loaded(
+                    drawable = ICON,
+                    contentDescription =
+                        ContentDescription.Resource(
+                            res = R.string.accessibility_wallet_button,
+                        ),
                 )
-            job.cancel()
-        }
+            )
+    }
 
     @Test
-    fun affordance_walletFeatureNotEnabled_modelIsNone() = runBlockingTest {
+    fun affordance_walletFeatureNotEnabled_modelIsNone() = testScope.runTest {
         setUpState(isWalletFeatureAvailable = false)
-        var latest: KeyguardQuickAffordanceConfig.LockScreenState? = null
 
-        val job = underTest.lockScreenState.onEach { latest = it }.launchIn(this)
+        val latest by collectLastValue(underTest.lockScreenState)
 
         assertThat(latest).isEqualTo(KeyguardQuickAffordanceConfig.LockScreenState.Hidden)
-
-        job.cancel()
     }
 
     @Test
-    fun affordance_queryNotSuccessful_modelIsNone() = runBlockingTest {
+    fun affordance_queryNotSuccessful_modelIsNone() = testScope.runTest {
         setUpState(isWalletQuerySuccessful = false)
-        var latest: KeyguardQuickAffordanceConfig.LockScreenState? = null
 
-        val job = underTest.lockScreenState.onEach { latest = it }.launchIn(this)
+        val latest by collectLastValue(underTest.lockScreenState)
 
         assertThat(latest).isEqualTo(KeyguardQuickAffordanceConfig.LockScreenState.Hidden)
-
-        job.cancel()
     }
 
     @Test
-    fun affordance_noSelectedCard_modelIsNone() = runBlockingTest {
+    fun affordance_noSelectedCard_modelIsNone() = testScope.runTest {
         setUpState(hasSelectedCard = false)
-        var latest: KeyguardQuickAffordanceConfig.LockScreenState? = null
 
-        val job = underTest.lockScreenState.onEach { latest = it }.launchIn(this)
+        val latest by collectLastValue(underTest.lockScreenState)
 
         assertThat(latest).isEqualTo(KeyguardQuickAffordanceConfig.LockScreenState.Hidden)
-
-        job.cancel()
     }
 
     @Test
@@ -179,7 +167,7 @@
     }
 
     @Test
-    fun getPickerScreenState_default() = runTest {
+    fun getPickerScreenState_default() = testScope.runTest {
         setUpState()
 
         assertThat(underTest.getPickerScreenState())
@@ -187,7 +175,7 @@
     }
 
     @Test
-    fun getPickerScreenState_unavailable() = runTest {
+    fun getPickerScreenState_unavailable() = testScope.runTest {
         setUpState(
             isWalletServiceAvailable = false,
         )
@@ -197,7 +185,7 @@
     }
 
     @Test
-    fun getPickerScreenState_disabledWhenTheFeatureIsNotEnabled() = runTest {
+    fun getPickerScreenState_disabledWhenTheFeatureIsNotEnabled() = testScope.runTest {
         setUpState(
             isWalletFeatureAvailable = false,
         )
@@ -207,7 +195,7 @@
     }
 
     @Test
-    fun getPickerScreenState_disabledWhenThereIsNoCard() = runTest {
+    fun getPickerScreenState_disabledWhenThereIsNoCard() = testScope.runTest {
         setUpState(
             hasSelectedCard = false,
         )
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/data/quickaffordance/QuickAccessWalletKeyguardQuickAffordanceConfig.kt b/packages/SystemUI/src/com/android/systemui/keyguard/data/quickaffordance/QuickAccessWalletKeyguardQuickAffordanceConfig.kt
index 7337292..a988a5c 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/data/quickaffordance/QuickAccessWalletKeyguardQuickAffordanceConfig.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/data/quickaffordance/QuickAccessWalletKeyguardQuickAffordanceConfig.kt
@@ -32,13 +32,19 @@
 import com.android.systemui.common.shared.model.Icon
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.plugins.ActivityStarter
 import com.android.systemui.wallet.controller.QuickAccessWalletController
 import com.android.systemui.wallet.util.getPaymentCards
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.ExperimentalCoroutinesApi
 import javax.inject.Inject
 import kotlinx.coroutines.channels.awaitClose
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.flatMapLatest
+import kotlinx.coroutines.flow.flowOf
 import kotlinx.coroutines.suspendCancellableCoroutine
+import kotlinx.coroutines.withContext
 
 /** Quick access wallet quick affordance data source. */
 @SysUISingleton
@@ -46,6 +52,7 @@
 @Inject
 constructor(
     @Application private val context: Context,
+    @Background private val backgroundDispatcher: CoroutineDispatcher,
     private val walletController: QuickAccessWalletController,
     private val activityStarter: ActivityStarter,
 ) : KeyguardQuickAffordanceConfig {
@@ -56,6 +63,7 @@
 
     override val pickerIconResourceId = R.drawable.ic_wallet_lockscreen
 
+    @OptIn(ExperimentalCoroutinesApi::class)
     override val lockScreenState: Flow<KeyguardQuickAffordanceConfig.LockScreenState> =
         conflatedCallbackFlow {
             val callback =
@@ -63,11 +71,7 @@
                     override fun onWalletCardsRetrieved(response: GetWalletCardsResponse) {
                         val hasCards = getPaymentCards(response.walletCards)?.isNotEmpty() == true
                         trySendWithFailureLogging(
-                            state(
-                                isFeatureEnabled = isWalletAvailable(),
-                                hasCard = hasCards,
-                                tileIcon = walletController.walletClient.tileIcon,
-                            ),
+                            hasCards,
                             TAG,
                         )
                     }
@@ -75,7 +79,7 @@
                     override fun onWalletCardRetrievalError(error: GetWalletCardsError) {
                         Log.e(TAG, "Wallet card retrieval error, message: \"${error?.message}\"")
                         trySendWithFailureLogging(
-                            KeyguardQuickAffordanceConfig.LockScreenState.Hidden,
+                            null,
                             TAG,
                         )
                     }
@@ -86,8 +90,12 @@
                 QuickAccessWalletController.WalletChangeEvent.WALLET_PREFERENCE_CHANGE,
                 QuickAccessWalletController.WalletChangeEvent.DEFAULT_PAYMENT_APP_CHANGE
             )
-            walletController.updateWalletPreference()
-            walletController.queryWalletCards(callback)
+
+            withContext(backgroundDispatcher) {
+                // Both must be called on background thread
+                walletController.updateWalletPreference()
+                walletController.queryWalletCards(callback)
+            }
 
             awaitClose {
                 walletController.unregisterWalletChangeObservers(
@@ -95,6 +103,19 @@
                     QuickAccessWalletController.WalletChangeEvent.DEFAULT_PAYMENT_APP_CHANGE
                 )
             }
+        }.flatMapLatest { hasCards ->
+            // If hasCards is null, this indicates an error occurred upon card retrieval
+            val state =
+                if (hasCards == null) {
+                    KeyguardQuickAffordanceConfig.LockScreenState.Hidden
+                } else {
+                    state(
+                        isWalletAvailable(),
+                        hasCards,
+                        walletController.walletClient.tileIcon,
+                    )
+                }
+            flowOf(state)
         }
 
     override suspend fun getPickerScreenState(): KeyguardQuickAffordanceConfig.PickerScreenState {
@@ -131,25 +152,33 @@
     }
 
     private suspend fun queryCards(): List<WalletCard> {
-        return suspendCancellableCoroutine { continuation ->
-            val callback =
-                object : QuickAccessWalletClient.OnWalletCardsRetrievedCallback {
-                    override fun onWalletCardsRetrieved(response: GetWalletCardsResponse) {
-                        continuation.resumeWith(
-                            Result.success(getPaymentCards(response.walletCards) ?: emptyList())
-                        )
-                    }
+        return withContext(backgroundDispatcher) {
+            suspendCancellableCoroutine { continuation ->
+                val callback =
+                    object : QuickAccessWalletClient.OnWalletCardsRetrievedCallback {
+                        override fun onWalletCardsRetrieved(response: GetWalletCardsResponse) {
+                            continuation.resumeWith(
+                                Result.success(getPaymentCards(response.walletCards))
+                            )
+                        }
 
-                    override fun onWalletCardRetrievalError(error: GetWalletCardsError) {
-                        continuation.resumeWith(Result.success(emptyList()))
+                        override fun onWalletCardRetrievalError(error: GetWalletCardsError) {
+                            continuation.resumeWith(Result.success(emptyList()))
+                        }
                     }
-                }
-            walletController.queryWalletCards(callback)
+                // Must be called on background thread
+                walletController.queryWalletCards(callback)
+            }
         }
     }
 
-    private fun isWalletAvailable() =
-        with(walletController.walletClient) { isWalletServiceAvailable && isWalletFeatureAvailable }
+    private suspend fun isWalletAvailable() =
+        withContext(backgroundDispatcher) {
+            with(walletController.walletClient) {
+                // Must be called on background thread
+                isWalletServiceAvailable && isWalletFeatureAvailable
+            }
+        }
 
     private fun state(
         isFeatureEnabled: Boolean,
diff --git a/packages/SystemUI/src/com/android/systemui/wallet/controller/QuickAccessWalletController.java b/packages/SystemUI/src/com/android/systemui/wallet/controller/QuickAccessWalletController.java
index 0e6df6b..e031be2 100644
--- a/packages/SystemUI/src/com/android/systemui/wallet/controller/QuickAccessWalletController.java
+++ b/packages/SystemUI/src/com/android/systemui/wallet/controller/QuickAccessWalletController.java
@@ -19,6 +19,7 @@
 import static com.android.systemui.wallet.controller.QuickAccessWalletController.WalletChangeEvent.DEFAULT_PAYMENT_APP_CHANGE;
 import static com.android.systemui.wallet.controller.QuickAccessWalletController.WalletChangeEvent.WALLET_PREFERENCE_CHANGE;
 
+import android.annotation.WorkerThread;
 import android.app.PendingIntent;
 import android.content.Context;
 import android.content.Intent;
@@ -146,7 +147,9 @@
 
     /**
      * Update the "show wallet" preference.
+     * This should not be called on the main thread.
      */
+    @WorkerThread
     public void updateWalletPreference() {
         mWalletEnabled = mQuickAccessWalletClient.isWalletServiceAvailable()
                 && mQuickAccessWalletClient.isWalletFeatureAvailable()
@@ -155,10 +158,12 @@
 
     /**
      * Query the wallet cards from {@link QuickAccessWalletClient}.
+     * This should not be called on the main thread.
      *
      * @param cardsRetriever a callback to retrieve wallet cards.
      * @param maxCards the maximum number of cards requested from the QuickAccessWallet
      */
+    @WorkerThread
     public void queryWalletCards(
             QuickAccessWalletClient.OnWalletCardsRetrievedCallback cardsRetriever, int maxCards) {
         if (mClock.elapsedRealtime() - mQawClientCreatedTimeMillis
@@ -182,9 +187,11 @@
 
     /**
      * Query the wallet cards from {@link QuickAccessWalletClient}.
+     * This should not be called on the main thread.
      *
      * @param cardsRetriever a callback to retrieve wallet cards.
      */
+    @WorkerThread
     public void queryWalletCards(
             QuickAccessWalletClient.OnWalletCardsRetrievedCallback cardsRetriever) {
         queryWalletCards(cardsRetriever, /* maxCards= */ 1);