Fix rotation issues with PaymentFlowActivity

stripe/src/main/java/com/stripe/android/view/PaymentFlowActivity.kt

@@ -47,9 +47,17 @@ class PaymentFlowActivity : StripeActivity() {
         paymentSessionData = requireNotNull(args.paymentSessionData) {
             "PaymentFlowActivity launched without PaymentSessionData"
         }
+        val paymentSessionConfig = args.paymentSessionConfig
+
+        val shippingInformation = savedInstanceState?.getParcelable(STATE_SHIPPING_INFO)
+            ?: paymentSessionConfig.prepopulatedShippingInfo
 
         paymentFlowPagerAdapter = PaymentFlowPagerAdapter(
-            this, args.paymentSessionConfig, customerSession
+            this,
+            paymentSessionConfig,
+            customerSession,
+            shippingInformation,
+            savedInstanceState?.getParcelable(STATE_SHIPPING_METHOD)
         )
         shipping_flow_viewpager.adapter = paymentFlowPagerAdapter
         shipping_flow_viewpager.addOnPageChangeListener(object : ViewPager.OnPageChangeListener {
@@ -100,6 +108,11 @@ class PaymentFlowActivity : StripeActivity() {
         title = paymentFlowPagerAdapter.getPageTitle(shipping_flow_viewpager.currentItem)
     }
 
+    override fun onRestoreInstanceState(savedInstanceState: Bundle) {
+        super.onRestoreInstanceState(savedInstanceState)
+        shipping_flow_viewpager.currentItem = savedInstanceState.getInt(STATE_CURRENT_ITEM, 0)
+    }
+
     public override fun onActionSave() {
         if (PaymentFlowPagerEnum.SHIPPING_INFO ==
             paymentFlowPagerAdapter.getPageAt(shipping_flow_viewpager.currentItem)) {
@@ -123,6 +136,13 @@ class PaymentFlowActivity : StripeActivity() {
         )
     }
 
+    override fun onSaveInstanceState(outState: Bundle) {
+        super.onSaveInstanceState(outState)
+        outState.putParcelable(STATE_SHIPPING_INFO, shippingInfo)
+        outState.putParcelable(STATE_SHIPPING_METHOD, selectedShippingMethod)
+        outState.putInt(STATE_CURRENT_ITEM, shipping_flow_viewpager.currentItem)
+    }
+
     @JvmSynthetic
     internal fun onShippingInfoSaved(shippingInformation: ShippingInformation?) {
         onShippingMethodsReady(validShippingMethods, defaultShippingMethod)
@@ -156,14 +176,31 @@ class PaymentFlowActivity : StripeActivity() {
     private fun onShippingInfoSubmitted() {
         hideKeyboard()
 
-        val shippingInfoWidget: ShippingInfoWidget = findViewById(R.id.shipping_info_widget)
-        shippingInfoWidget.shippingInformation?.let { shippingInformation ->
-            shippingInformationSubmitted = shippingInformation
+        shippingInfo?.let { shippingInfo ->
+            shippingInformationSubmitted = shippingInfo
             setCommunicatingProgress(true)
-            broadcastShippingInfoSubmitted(shippingInformation)
+            broadcastShippingInfoSubmitted(shippingInfo)
         }
     }
 
+    private val shippingInfo: ShippingInformation?
+        get() {
+            val shippingInfoWidget: ShippingInfoWidget = findViewById(R.id.shipping_info_widget)
+            return shippingInfoWidget.rawShippingInformation
+        }
+
+    private val selectedShippingMethod: ShippingMethod?
+        get() {
+            return if (PaymentFlowPagerEnum.SHIPPING_METHOD ==
+                paymentFlowPagerAdapter.getPageAt(shipping_flow_viewpager.currentItem)) {
+                val selectShippingMethodWidget: SelectShippingMethodWidget =
+                    findViewById(R.id.select_shipping_method_widget)
+                selectShippingMethodWidget.selectedShippingMethod
+            } else {
+                null
+            }
+        }
+
     private fun hideKeyboard() {
         val inputMethodManager: InputMethodManager =
             getSystemService(Context.INPUT_METHOD_SERVICE) as InputMethodManager
@@ -203,9 +240,9 @@ class PaymentFlowActivity : StripeActivity() {
     override fun onBackPressed() {
         if (hasPreviousPage()) {
             shipping_flow_viewpager.currentItem = shipping_flow_viewpager.currentItem - 1
-            return
+        } else {
+            super.onBackPressed()
         }
-        super.onBackPressed()
     }
 
     private class CustomerShippingInfoSavedListener internal constructor(
@@ -226,5 +263,9 @@ class PaymentFlowActivity : StripeActivity() {
         internal const val TOKEN_PAYMENT_FLOW_ACTIVITY: String = "PaymentFlowActivity"
         internal const val TOKEN_SHIPPING_INFO_SCREEN: String = "ShippingInfoScreen"
         internal const val TOKEN_SHIPPING_METHOD_SCREEN: String = "ShippingMethodScreen"
+
+        private const val STATE_SHIPPING_INFO: String = "state_shipping_info"
+        private const val STATE_SHIPPING_METHOD: String = "state_shipping_method"
+        private const val STATE_CURRENT_ITEM: String = "state_current_item"
     }
 }

stripe/src/main/java/com/stripe/android/view/PaymentFlowPagerAdapter.kt

@@ -1,20 +1,25 @@
 package com.stripe.android.view
 
 import android.content.Context
+import android.os.Parcelable
 import android.view.LayoutInflater
 import android.view.View
 import android.view.ViewGroup
 import androidx.viewpager.widget.PagerAdapter
 import com.stripe.android.CustomerSession
 import com.stripe.android.PaymentSessionConfig
 import com.stripe.android.R
+import com.stripe.android.model.ShippingInformation
 import com.stripe.android.model.ShippingMethod
 import java.util.ArrayList
+import kotlinx.android.parcel.Parcelize
 
 internal class PaymentFlowPagerAdapter(
     private val context: Context,
     private val paymentSessionConfig: PaymentSessionConfig,
-    private val customerSession: CustomerSession
+    private val customerSession: CustomerSession,
+    private val shippingInformation: ShippingInformation?,
+    private val shippingMethod: ShippingMethod?
 ) : PagerAdapter() {
     private val pages: MutableList<PaymentFlowPagerEnum>
 
@@ -61,25 +66,33 @@ internal class PaymentFlowPagerAdapter(
 
     override fun instantiateItem(collection: ViewGroup, position: Int): Any {
         val paymentFlowPagerEnum = pages[position]
-        val inflater = LayoutInflater.from(context)
-        val layout = inflater
+        val layout = LayoutInflater.from(context)
             .inflate(paymentFlowPagerEnum.layoutResId, collection, false) as ViewGroup
-        if (paymentFlowPagerEnum == PaymentFlowPagerEnum.SHIPPING_METHOD) {
-            customerSession
-                .addProductUsageTokenIfValid(PaymentFlowActivity.TOKEN_SHIPPING_METHOD_SCREEN)
-            val selectShippingMethodWidget =
-                layout.findViewById<SelectShippingMethodWidget>(R.id.select_shipping_method_widget)
-            selectShippingMethodWidget
-                .setShippingMethods(validShippingMethods, defaultShippingMethod)
-        }
-        if (paymentFlowPagerEnum == PaymentFlowPagerEnum.SHIPPING_INFO) {
-            customerSession
-                .addProductUsageTokenIfValid(PaymentFlowActivity.TOKEN_SHIPPING_INFO_SCREEN)
-            val shippingInfoWidget =
-                layout.findViewById<ShippingInfoWidget>(R.id.shipping_info_widget)
-            shippingInfoWidget.setHiddenFields(paymentSessionConfig.hiddenShippingInfoFields)
-            shippingInfoWidget.setOptionalFields(paymentSessionConfig.optionalShippingInfoFields)
-            shippingInfoWidget.populateShippingInfo(paymentSessionConfig.prepopulatedShippingInfo)
+
+        when (paymentFlowPagerEnum) {
+            PaymentFlowPagerEnum.SHIPPING_METHOD -> {
+                customerSession
+                    .addProductUsageTokenIfValid(PaymentFlowActivity.TOKEN_SHIPPING_METHOD_SCREEN)
+                val selectShippingMethodWidget: SelectShippingMethodWidget =
+                    layout.findViewById(R.id.select_shipping_method_widget)
+                selectShippingMethodWidget
+                    .setShippingMethods(validShippingMethods, defaultShippingMethod)
+                shippingMethod?.let {
+                    selectShippingMethodWidget.setSelectedShippingMethod(it)
+                }
+            }
+            PaymentFlowPagerEnum.SHIPPING_INFO -> {
+                customerSession
+                    .addProductUsageTokenIfValid(PaymentFlowActivity.TOKEN_SHIPPING_INFO_SCREEN)
+                val shippingInfoWidget: ShippingInfoWidget =
+                    layout.findViewById(R.id.shipping_info_widget)
+                shippingInfoWidget
+                    .setHiddenFields(paymentSessionConfig.hiddenShippingInfoFields)
+                shippingInfoWidget
+                    .setOptionalFields(paymentSessionConfig.optionalShippingInfoFields)
+                shippingInfoWidget
+                    .populateShippingInfo(shippingInformation)
+            }
         }
         collection.addView(layout)
         return layout
@@ -93,10 +106,8 @@ internal class PaymentFlowPagerAdapter(
         return pages.size
     }
 
-    fun getPageAt(position: Int): PaymentFlowPagerEnum? {
-        return if (position < pages.size) {
-            pages[position]
-        } else null
+    internal fun getPageAt(position: Int): PaymentFlowPagerEnum? {
+        return pages.getOrNull(position)
     }
 
     override fun isViewFromObject(view: View, o: Any): Boolean {
@@ -106,4 +117,28 @@ internal class PaymentFlowPagerAdapter(
     override fun getPageTitle(position: Int): CharSequence? {
         return context.getString(pages[position].titleResId)
     }
+
+    override fun saveState(): Parcelable? {
+        return State(pages, shippingInfoSaved, validShippingMethods, defaultShippingMethod)
+    }
+
+    override fun restoreState(state: Parcelable?, loader: ClassLoader?) {
+        if (state is State) {
+            this.pages.clear()
+            this.pages.addAll(state.pages)
+            this.shippingInfoSaved = state.shippingInfoSaved
+            this.validShippingMethods = state.validShippingMethods
+            this.defaultShippingMethod = state.defaultShippingMethod
+
+            notifyDataSetChanged()
+        }
+    }
+
+    @Parcelize
+    internal class State(
+        internal val pages: List<PaymentFlowPagerEnum>,
+        internal val shippingInfoSaved: Boolean,
+        internal val validShippingMethods: List<ShippingMethod>?,
+        internal val defaultShippingMethod: ShippingMethod?
+    ) : Parcelable
 }

stripe/src/main/java/com/stripe/android/view/SelectShippingMethodWidget.kt

@@ -43,4 +43,8 @@ internal class SelectShippingMethodWidget @JvmOverloads constructor(
     ) {
         shippingMethodAdapter.setShippingMethods(shippingMethods, defaultShippingMethod)
     }
+
+    fun setSelectedShippingMethod(shippingMethod: ShippingMethod) {
+        shippingMethodAdapter.setSelected(shippingMethod)
+    }
 }

stripe/src/main/java/com/stripe/android/view/ShippingInfoWidget.kt

@@ -42,12 +42,23 @@ class ShippingInfoWidget @JvmOverloads constructor(
     private val stateEditText: StripeEditText
     private val phoneNumberEditText: StripeEditText
 
+    /**
+     * Return [ShippingInformation] based on user input if valid, otherwise null.
+     */
     val shippingInformation: ShippingInformation?
         get() {
-            if (!validateAllFields()) {
-                return null
+            return if (!validateAllFields()) {
+                null
+            } else {
+                rawShippingInformation
             }
+        }
 
+    /**
+     * Return [ShippingInformation] based on user input.
+     */
+    internal val rawShippingInformation: ShippingInformation
+        get() {
             val address = Address.Builder()
                 .setCity(cityEditText.text?.toString())
                 .setCountry(countryAutoCompleteTextView.selectedCountryCode)
@@ -151,11 +162,12 @@ class ShippingInfoWidget @JvmOverloads constructor(
             return
         }
 
-        val address = shippingInformation.address
-        if (address != null) {
+        shippingInformation.address?.let { address ->
             cityEditText.setText(address.city)
-            if (address.country != null && address.country.isNotEmpty()) {
-                countryAutoCompleteTextView.setCountrySelected(address.country)
+            address.country?.let { country ->
+                if (country.isNotEmpty()) {
+                    countryAutoCompleteTextView.setCountrySelected(country)
+                }
             }
             addressEditText.setText(address.line1)
             addressEditText2.setText(address.line2)

stripe/src/main/java/com/stripe/android/view/ShippingMethodAdapter.kt

@@ -59,11 +59,19 @@ internal class ShippingMethodAdapter :
         notifyItemChanged(selectedIndex)
     }
 
+    internal fun setSelected(shippingMethod: ShippingMethod) {
+        val previouslySelectedIndex = selectedIndex
+        selectedIndex = shippingMethods.indexOf(shippingMethod)
+        if (previouslySelectedIndex != selectedIndex) {
+            notifyItemChanged(previouslySelectedIndex)
+            notifyItemChanged(selectedIndex)
+        }
+    }
+
     internal class ShippingMethodViewHolder constructor(
         private val shippingMethodView: ShippingMethodView,
         adapter: ShippingMethodAdapter
     ) : RecyclerView.ViewHolder(shippingMethodView) {
-
         init {
             shippingMethodView.setOnClickListener {
                 adapter.onShippingMethodSelected(adapterPosition)