Add a selection holder for embedded. (#9791)

paymentsheet/src/main/java/com/stripe/android/paymentelement/embedded/EmbeddedSelectionHolder.kt

@@ -0,0 +1,22 @@
+package com.stripe.android.paymentelement.embedded
+
+import androidx.lifecycle.SavedStateHandle
+import com.stripe.android.paymentsheet.model.PaymentSelection
+import kotlinx.coroutines.flow.StateFlow
+import javax.inject.Inject
+import javax.inject.Singleton
+
+@Singleton
+internal class EmbeddedSelectionHolder @Inject constructor(
+    private val savedStateHandle: SavedStateHandle,
+) {
+    val selection: StateFlow<PaymentSelection?> = savedStateHandle.getStateFlow(EMBEDDED_SELECTION_KEY, null)
+
+    fun set(updatedSelection: PaymentSelection?) {
+        savedStateHandle[EMBEDDED_SELECTION_KEY] = updatedSelection
+    }
+
+    companion object {
+        const val EMBEDDED_SELECTION_KEY = "EMBEDDED_SELECTION_KEY"
+    }
+}

paymentsheet/src/main/java/com/stripe/android/paymentelement/embedded/SharedPaymentElementViewModel.kt

@@ -65,6 +65,7 @@ import kotlinx.coroutines.Dispatchers
 import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
 import kotlinx.coroutines.flow.asStateFlow
+import kotlinx.coroutines.launch
 import kotlinx.coroutines.plus
 import javax.inject.Inject
 import javax.inject.Named
@@ -79,6 +80,7 @@ internal class SharedPaymentElementViewModel @Inject constructor(
     @IOContext ioContext: CoroutineContext,
     private val configurationHandler: EmbeddedConfigurationHandler,
     private val paymentOptionDisplayDataFactory: PaymentOptionDisplayDataFactory,
+    private val selectionHolder: EmbeddedSelectionHolder,
 ) : ViewModel() {
     private val _paymentOption: MutableStateFlow<PaymentOptionDisplayData?> = MutableStateFlow(null)
     val paymentOption: StateFlow<PaymentOptionDisplayData?> = _paymentOption.asStateFlow()
@@ -91,6 +93,15 @@ internal class SharedPaymentElementViewModel @Inject constructor(
     @Volatile
     var confirmationState: EmbeddedConfirmationHelper.State? = null
 
+    init {
+        viewModelScope.launch {
+            selectionHolder.selection.collect { selection ->
+                _paymentOption.value = paymentOptionDisplayDataFactory.create(selection)
+                confirmationState = confirmationState?.copy(selection = selection)
+            }
+        }
+    }
+
     suspend fun configure(
         intentConfiguration: PaymentSheet.IntentConfiguration,
         configuration: EmbeddedPaymentElement.Configuration,
@@ -106,7 +117,7 @@ internal class SharedPaymentElementViewModel @Inject constructor(
                     initializationMode = PaymentElementLoader.InitializationMode.DeferredIntent(intentConfiguration),
                     configuration = configuration,
                 )
-                _paymentOption.value = paymentOptionDisplayDataFactory.create(state.paymentSelection)
+                selectionHolder.set(state.paymentSelection)
                 ConfigureResult.Succeeded()
             },
             onFailure = { error ->

paymentsheet/src/test/java/com/stripe/android/paymentelement/embedded/EmbeddedSelectionHolderTest.kt

@@ -0,0 +1,59 @@
+package com.stripe.android.paymentelement.embedded
+
+import androidx.lifecycle.SavedStateHandle
+import app.cash.turbine.test
+import com.google.common.truth.Truth.assertThat
+import com.stripe.android.paymentsheet.model.PaymentSelection
+import com.stripe.android.paymentsheet.model.paymentMethodType
+import kotlinx.coroutines.test.runTest
+import kotlin.test.Test
+
+internal class EmbeddedSelectionHolderTest {
+    @Test
+    fun `setting selection emits value in selection state flow`() = testScenario {
+        selectionHolder.selection.test {
+            assertThat(awaitItem()).isNull()
+            selectionHolder.set(PaymentSelection.GooglePay)
+            assertThat(awaitItem()?.paymentMethodType).isEqualTo("google_pay")
+        }
+    }
+
+    @Test
+    fun `setting selection updates savedStateHandle`() = testScenario {
+        assertThat(savedStateHandle.get<PaymentSelection?>(EmbeddedSelectionHolder.EMBEDDED_SELECTION_KEY))
+            .isNull()
+        selectionHolder.set(PaymentSelection.GooglePay)
+        assertThat(savedStateHandle.get<PaymentSelection?>(EmbeddedSelectionHolder.EMBEDDED_SELECTION_KEY))
+            .isEqualTo(PaymentSelection.GooglePay)
+    }
+
+    @Test
+    fun `initializing with selection in savedStateHandle sets initial value`() = testScenario(
+        setup = {
+            set(EmbeddedSelectionHolder.EMBEDDED_SELECTION_KEY, PaymentSelection.GooglePay)
+        },
+    ) {
+        assertThat(savedStateHandle.get<PaymentSelection?>(EmbeddedSelectionHolder.EMBEDDED_SELECTION_KEY))
+            .isEqualTo(PaymentSelection.GooglePay)
+        selectionHolder.set(null)
+        assertThat(savedStateHandle.get<PaymentSelection?>(EmbeddedSelectionHolder.EMBEDDED_SELECTION_KEY))
+            .isNull()
+    }
+
+    private class Scenario(
+        val selectionHolder: EmbeddedSelectionHolder,
+        val savedStateHandle: SavedStateHandle,
+    )
+
+    private fun testScenario(
+        setup: SavedStateHandle.() -> Unit = {},
+        block: suspend Scenario.() -> Unit,
+    ) = runTest {
+        val savedStateHandle = SavedStateHandle()
+        setup(savedStateHandle)
+        Scenario(
+            selectionHolder = EmbeddedSelectionHolder(savedStateHandle),
+            savedStateHandle = savedStateHandle,
+        ).block()
+    }
+}

paymentsheet/src/test/java/com/stripe/android/paymentelement/embedded/SharedPaymentElementViewModelTest.kt

@@ -1,5 +1,6 @@
 package com.stripe.android.paymentelement.embedded
 
+import androidx.lifecycle.SavedStateHandle
 import androidx.test.core.app.ApplicationProvider
 import app.cash.turbine.Turbine
 import app.cash.turbine.test
@@ -13,6 +14,7 @@ import com.stripe.android.paymentelement.ExperimentalEmbeddedPaymentElementApi
 import com.stripe.android.paymentelement.confirmation.FakeConfirmationHandler
 import com.stripe.android.paymentsheet.PaymentSheet
 import com.stripe.android.paymentsheet.model.PaymentSelection
+import com.stripe.android.paymentsheet.model.paymentMethodType
 import com.stripe.android.paymentsheet.state.PaymentElementLoader
 import com.stripe.android.ui.core.cbc.CardBrandChoiceEligibility
 import kotlinx.coroutines.test.runTest
@@ -64,6 +66,47 @@ internal class SharedPaymentElementViewModelTest {
         assertThat(viewModel.confirmationState?.paymentMethodMetadata).isNotNull()
     }
 
+    @Test
+    fun `Updating selection updates confirmationState`() = testScenario {
+        val configuration = EmbeddedPaymentElement.Configuration.Builder("Example, Inc.").build()
+        configurationHandler.emit(
+            Result.success(
+                PaymentElementLoader.State(
+                    config = configuration.asCommonConfiguration(),
+                    customer = null,
+                    linkState = null,
+                    paymentSelection = PaymentSelection.GooglePay,
+                    validationError = null,
+                    paymentMethodMetadata = PaymentMethodMetadataFactory.create(
+                        stripeIntent = PaymentIntentFixtures.PI_SUCCEEDED,
+                        billingDetailsCollectionConfiguration = configuration
+                            .billingDetailsCollectionConfiguration,
+                        allowsDelayedPaymentMethods = configuration.allowsDelayedPaymentMethods,
+                        allowsPaymentMethodsRequiringShippingAddress = configuration
+                            .allowsPaymentMethodsRequiringShippingAddress,
+                        isGooglePayReady = true,
+                        cbcEligibility = CardBrandChoiceEligibility.Ineligible,
+                    ),
+                )
+            )
+        )
+
+        assertThat(viewModel.confirmationState?.selection?.paymentMethodType).isNull()
+
+        assertThat(
+            viewModel.configure(
+                PaymentSheet.IntentConfiguration(
+                    PaymentSheet.IntentConfiguration.Mode.Payment(5000, "USD"),
+                ),
+                configuration = configuration,
+            )
+        ).isInstanceOf<EmbeddedPaymentElement.ConfigureResult.Succeeded>()
+
+        assertThat(viewModel.confirmationState?.selection?.paymentMethodType).isEqualTo("google_pay")
+        selectionHolder.set(null)
+        assertThat(viewModel.confirmationState?.selection?.paymentMethodType).isNull()
+    }
+
     @Test
     fun `configure maps success result`() = testScenario {
         val configuration = EmbeddedPaymentElement.Configuration.Builder("Example, Inc.").build()
@@ -103,7 +146,7 @@ internal class SharedPaymentElementViewModelTest {
     }
 
     @Test
-    fun `configure emits payment option`() = testScenario {
+    fun `configure emits payment option and sets initial selection`() = testScenario {
         val configuration = EmbeddedPaymentElement.Configuration.Builder("Example, Inc.").build()
         configurationHandler.emit(
             Result.success(
@@ -127,6 +170,7 @@ internal class SharedPaymentElementViewModelTest {
             )
         )
 
+        assertThat(selectionHolder.selection.value?.paymentMethodType).isNull()
         viewModel.paymentOption.test {
             assertThat(awaitItem()).isNull()
 
@@ -141,6 +185,7 @@ internal class SharedPaymentElementViewModelTest {
 
             assertThat(awaitItem()?.paymentMethodType).isEqualTo("google_pay")
         }
+        assertThat(selectionHolder.selection.value?.paymentMethodType).isEqualTo("google_pay")
     }
 
     @Test
@@ -166,18 +211,21 @@ internal class SharedPaymentElementViewModelTest {
             iconLoader = mock(),
             context = ApplicationProvider.getApplicationContext(),
         )
+        val selectionHolder = EmbeddedSelectionHolder(SavedStateHandle())
 
         runTest {
             val viewModel = SharedPaymentElementViewModel(
                 confirmationHandlerFactory = { confirmationHandler },
                 ioContext = testScheduler,
                 configurationHandler = configurationHandler,
                 paymentOptionDisplayDataFactory = paymentOptionDisplayDataFactory,
+                selectionHolder = selectionHolder,
             )
 
             Scenario(
                 configurationHandler = configurationHandler,
                 viewModel = viewModel,
+                selectionHolder = selectionHolder,
             ).block()
         }
 
@@ -188,6 +236,7 @@ internal class SharedPaymentElementViewModelTest {
     private class Scenario(
         val configurationHandler: FakeEmbeddedConfigurationHandler,
         val viewModel: SharedPaymentElementViewModel,
+        val selectionHolder: EmbeddedSelectionHolder,
     )
 
     private class FakeEmbeddedConfigurationHandler : EmbeddedConfigurationHandler {