Fix keeping row selection across updates even when form is invalid

Example/PaymentSheet Example/PaymentSheetUITest/EmbeddedUITest.swift

@@ -50,6 +50,7 @@ class EmbeddedUITests: PaymentSheetUITestCase {
         XCTAssertTrue(app.buttons["Reload"].waitForExistence(timeout: 10))
         // ...should cause Card to no longer be the selected payment method.
         XCTAssertFalse(app.staticTexts["Payment method"].exists)
+        XCTAssertFalse(app.buttons["Card"].isSelected)
 
         // ....Tapping card should show the card form with details preserved
         app.buttons["Card"].waitForExistenceAndTap()
@@ -65,6 +66,7 @@ class EmbeddedUITests: PaymentSheetUITestCase {
         // ...card entered for setup should be preserved after update
         XCTAssertTrue(app.staticTexts["Payment method"].waitForExistence(timeout: 10))
         XCTAssertEqual(app.staticTexts["Payment method"].label, "•••• 4242")
+        XCTAssertTrue(app.buttons["Card"].isSelected)
 
         // ...selecting Alipay w/ deferred PaymentIntent...
         app.buttons["Alipay"].waitForExistenceAndTap()
@@ -82,6 +84,7 @@ class EmbeddedUITests: PaymentSheetUITestCase {
         XCTAssertTrue(app.buttons["Reload"].waitForExistence(timeout: 10))
         // ...should cause Alipay to no longer be the selected payment method, since it is not valid for setup.
         XCTAssertFalse(app.staticTexts["Payment method"].exists)
+        XCTAssertFalse(app.buttons["Alipay"].exists)
 
         // ...go back into deferred PaymentIntent mode
         app.buttons.matching(identifier: "Payment").element(boundBy: 1).waitForExistenceAndTap()
@@ -96,6 +99,7 @@ class EmbeddedUITests: PaymentSheetUITestCase {
         XCTAssertTrue(app.buttons["Reload"].waitForExistence(timeout: 10))
         // ...should cause Cash App Pay to be the selected payment method, since it is valid for setup.
         XCTAssertEqual(app.staticTexts["Payment method"].label, "Cash App Pay")
+        XCTAssertTrue(app.buttons["Cash App Pay"].isSelected)
 
         // ...go back into deferred PaymentIntent mode
         app.buttons.matching(identifier: "Payment").element(boundBy: 1).waitForExistenceAndTap()
@@ -115,9 +119,11 @@ class EmbeddedUITests: PaymentSheetUITestCase {
         XCTAssertTrue(app.buttons["Reload"].waitForExistence(timeout: 10))
         // ...should cause Klarna to no longer be the selected payment method.
         XCTAssertFalse(app.staticTexts["Payment method"].exists)
+        XCTAssertFalse(app.buttons["Klarna"].isSelected)
         // ...selecting Klarna should present a Klarna form with the previously entered email
         app.buttons["Klarna"].waitForExistenceAndTap()
         app.buttons["Continue"].waitForExistenceAndTap()
+        XCTAssertTrue(app.buttons["Klarna"].isSelected)
 
         let klarnaAnalytics = analyticsLog.compactMap({ $0[string: "event"] })
         XCTAssertEqual(
@@ -131,6 +137,7 @@ class EmbeddedUITests: PaymentSheetUITestCase {
         XCTAssertTrue(app.buttons["Reload"].waitForExistence(timeout: 10))
         // ... Klarna should still be selected
         XCTAssertEqual(app.staticTexts["Payment method"].label, "Klarna")
+        XCTAssertTrue(app.buttons["Klarna"].isSelected)
 
         // Confirm the Klarna payment
         XCTAssertTrue(app.buttons["Checkout"].waitForExistence(timeout: 10))

StripePaymentSheet/StripePaymentSheet/Source/PaymentSheet/Embedded/EmbeddedPaymentElement+Internal.swift

@@ -16,7 +16,7 @@ extension EmbeddedPaymentElement {
         configuration: Configuration,
         loadResult: PaymentSheetLoader.LoadResult,
         analyticsHelper: PaymentSheetAnalyticsHelper,
-        previousPaymentOption: PaymentOption? = nil,
+        previousSelection: RowButtonType? = nil,
         delegate: EmbeddedPaymentMethodsViewDelegate? = nil
     ) -> EmbeddedPaymentMethodsView {
         // Restore the customer's previous payment method.
@@ -35,20 +35,9 @@ extension EmbeddedPaymentElement {
             isFlatCheckmarkStyle: configuration.appearance.embeddedPaymentElement.row.style == .flatWithCheckmark
         )
         let initialSelection: RowButtonType? = {
-            // Select the previous payment option
-            switch previousPaymentOption {
-            case .applePay:
-                return .applePay
-            case .link:
-                return .link
-            case .external(paymentMethod: let paymentMethod, billingDetails: _):
-                return .new(paymentMethodType: .external(paymentMethod))
-            case .saved(paymentMethod: let paymentMethod, confirmParams: _):
-                return .saved(paymentMethod: paymentMethod)
-            case .new(confirmParams: let confirmParams):
-                return .new(paymentMethodType: confirmParams.paymentMethodType)
-            case nil:
-                break
+            // First, respect the previous selection
+            if let previousSelection {
+                return previousSelection
             }
 
             // If there's no previous customer input, default to the customer's default or the first saved payment method, if any
@@ -103,9 +92,16 @@ extension EmbeddedPaymentElement {
     }
 
     // Helper method to create Form VC for a payment method row, if applicable.
-    func makeFormViewControllerIfNecessary(
+    static func makeFormViewControllerIfNecessary(
         selection: RowButtonType?,
-        previousPaymentOption: PaymentOption?
+        previousPaymentOption: PaymentOption?,
+        configuration: Configuration,
+        intent: Intent,
+        elementsSession: STPElementsSession,
+        savedPaymentMethods: [STPPaymentMethod],
+        analyticsHelper: PaymentSheetAnalyticsHelper,
+        formCache: PaymentMethodFormCache,
+        delegate: EmbeddedFormViewControllerDelegate
     ) -> EmbeddedFormViewController? {
         guard case let .new(paymentMethodType) = selection else {
             return nil
@@ -117,10 +113,10 @@ extension EmbeddedPaymentElement {
             elementsSession: elementsSession,
             shouldUseNewCardNewCardHeader: savedPaymentMethods.first?.type == .card,
             paymentMethodType: paymentMethodType,
-            previousPaymentOption:previousPaymentOption,
+            previousPaymentOption: previousPaymentOption,
             analyticsHelper: analyticsHelper,
             formCache: formCache,
-            delegate: self
+            delegate: delegate
         )
         guard formViewController.collectsUserInput else {
             return nil
@@ -139,9 +135,16 @@ extension EmbeddedPaymentElement: EmbeddedPaymentMethodsViewDelegate {
     func embeddedPaymentMethodsViewDidUpdateSelection() {
         // 1. Update the currently selection's form VC to match the selection.
         // Note `paymentOption` derives from this property
-        self.selectedFormViewController = makeFormViewControllerIfNecessary(
+        self.selectedFormViewController = Self.makeFormViewControllerIfNecessary(
             selection: embeddedPaymentMethodsView.selectedRowButton?.type,
-            previousPaymentOption:  selectedFormViewController?.previousPaymentOption
+            previousPaymentOption:  selectedFormViewController?.previousPaymentOption,
+            configuration: configuration,
+            intent: intent,
+            elementsSession: elementsSession,
+            savedPaymentMethods: savedPaymentMethods,
+            analyticsHelper: analyticsHelper,
+            formCache: formCache,
+            delegate: self
         )
 
         // 2. Inform the delegate of the updated payment option

StripePaymentSheet/StripePaymentSheet/Source/PaymentSheet/Embedded/EmbeddedPaymentElement.swift

@@ -118,7 +118,7 @@ public final class EmbeddedPaymentElement {
         latestUpdateTask?.cancel()
         _ = await latestUpdateTask?.value
         // Start the new update task
-        let currentUpdateTask = Task { @MainActor [weak self, configuration, analyticsHelper] in
+        let currentUpdateTask: Task<UpdateResult, Never> = Task { @MainActor [weak self, configuration, analyticsHelper] in
             // ⚠️ Don't modify `self` until after all `awaits` to avoid being canceled halfway through and leaving self in a partially updated state.
             // 1. Reload v1/elements/session.
             let loadResult: PaymentSheetLoader.LoadResult
@@ -142,17 +142,52 @@ public final class EmbeddedPaymentElement {
             self.loadResult = loadResult
             self.savedPaymentMethods = loadResult.savedPaymentMethods
             self.formCache = .init() // Clear the cache because the form may have changed e.g. different mandate or different fields.
+            let isPreviousPaymentOptionStillDisplayed: Bool = {
+                switch previousPaymentOption {
+                case .none:
+                    return true
+                case .applePay:
+                    return PaymentSheet.isApplePayEnabled(elementsSession: loadResult.elementsSession, configuration: configuration)
+                case .link:
+                    return PaymentSheet.isLinkEnabled(elementsSession: loadResult.elementsSession, configuration: configuration)
+                case .saved(paymentMethod: let paymentMethod, confirmParams: _):
+                    return loadResult.savedPaymentMethods.contains(paymentMethod)
+                case .new(confirmParams: let confirmParams):
+                    return loadResult.paymentMethodTypes.contains(confirmParams.paymentMethodType)
+                case .external(paymentMethod: let paymentMethod, billingDetails: _):
+                    return loadResult.paymentMethodTypes.contains(.external(paymentMethod))
+                }
+            }()
+            let previousSelectedRowType = self.embeddedPaymentMethodsView.selectedRowButton?.type
+            // Make the new form VC for the previously selected row type if it's still in the list
+            let selectedFormViewController = Self.makeFormViewControllerIfNecessary(
+                selection: isPreviousPaymentOptionStillDisplayed ? previousSelectedRowType : nil,
+                previousPaymentOption: previousPaymentOption,
+                configuration: self.configuration,
+                intent: loadResult.intent,
+                elementsSession: loadResult.elementsSession,
+                savedPaymentMethods: loadResult.savedPaymentMethods,
+                analyticsHelper: self.analyticsHelper,
+                formCache: self.formCache,
+                delegate: self
+            )
+            self.selectedFormViewController = selectedFormViewController
+            // Make the new list view, selecting the previous row if it's still in the list and it doesn't have a form or it's form is valid
+            let shouldSelectPreviousRow: Bool = {
+                guard isPreviousPaymentOptionStillDisplayed else { return false }
+                if let selectedFormViewController {
+                    return selectedFormViewController.selectedPaymentOption != nil
+                } else {
+                    return true
+                }
+            }()
             self.embeddedPaymentMethodsView = Self.makeView(
                 configuration: configuration,
                 loadResult: loadResult,
                 analyticsHelper: analyticsHelper,
-                previousPaymentOption: previousPaymentOption,
+                previousSelection: shouldSelectPreviousRow ? previousSelectedRowType : nil,
                 delegate: self
             )
-            self.selectedFormViewController = makeFormViewControllerIfNecessary(
-                selection: self.embeddedPaymentMethodsView.selectedRowButton?.type,
-                previousPaymentOption: previousPaymentOption
-            )
             self.containerView.updateEmbeddedPaymentMethodsView(embeddedPaymentMethodsView)
             informDelegateIfPaymentOptionUpdated()
             return .succeeded