Add support for retrieving gateway-specific saved tokens

changelog/fix-token-retrieval-per-payment-method

@@ -0,0 +1,4 @@
+Significance: minor
+Type: fix
+
+Retrieve saved tokens only relevant for the specific payment gateway.

includes/class-wc-payments-token-service.php

@@ -163,15 +163,7 @@ public function woocommerce_get_customer_payment_tokens( $tokens, $user_id, $gat
 				}
 			}
 
-			$retrievable_payment_method_types = [ Payment_Method::CARD ];
-
-			if ( in_array( Payment_Method::SEPA, WC_Payments::get_gateway()->get_upe_enabled_payment_method_ids(), true ) ) {
-				$retrievable_payment_method_types[] = Payment_Method::SEPA;
-			}
-
-			if ( in_array( Payment_Method::LINK, WC_Payments::get_gateway()->get_upe_enabled_payment_method_ids(), true ) ) {
-				$retrievable_payment_method_types[] = Payment_Method::LINK;
-			}
+			$retrievable_payment_method_types = $this->get_retrievable_payment_method_types( $gateway_id );
 
 			$payment_methods = [];
 
@@ -213,6 +205,77 @@ public function woocommerce_get_customer_payment_tokens( $tokens, $user_id, $gat
 		return $tokens;
 	}
 
+	/**
+	 * Retrieves the payment method types for which tokens should be retrieved.
+	 *
+	 * This function determines the appropriate payment method types based on the provided gateway ID.
+	 * - If a gateway ID is provided, it retrieves the payment methods specific to that gateway to prevent duplication of saved tokens under incorrect payment methods during checkout.
+	 * - If no gateway ID is provided, it retrieves the default payment methods to fetch all saved tokens, e.g., for the Blocks checkout or My Account page.
+	 *
+	 * @param string|null $gateway_id The optional ID of the gateway.
+	 * @return array The list of retrievable payment method types.
+	 */
+	private function get_retrievable_payment_method_types( $gateway_id = null ) {
+		if ( empty( $gateway_id ) ) {
+			return $this->get_all_retrievable_payment_types();
+		} else {
+			return $this->get_gateway_specific_retrievable_payment_types( $gateway_id );
+		}
+	}
+
+	/**
+	 * Returns all the enabled retrievable payment method types.
+	 *
+	 * @return array Enabled retrievable payment method types.
+	 */
+	private function get_all_retrievable_payment_types() {
+		$types = [ Payment_Method::CARD ];
+
+		if ( $this->is_payment_method_enabled( Payment_Method::SEPA ) ) {
+			$types[] = Payment_Method::SEPA;
+		}
+
+		if ( $this->is_payment_method_enabled( Payment_Method::LINK ) ) {
+			$types[] = Payment_Method::LINK;
+		}
+
+		return $types;
+	}
+	/**
+	 * Returns retrievable payment method types for a given gateway.
+	 *
+	 * @param string $gateway_id The ID of the gateway.
+	 * @return array Retrievable payment method types for the specified gateway.
+	 */
+	private function get_gateway_specific_retrievable_payment_types( $gateway_id ) {
+		$types = [];
+
+		foreach ( self::REUSABLE_GATEWAYS_BY_PAYMENT_METHOD as $payment_method => $gateway ) {
+			if ( $gateway !== $gateway_id ) {
+				continue;
+			}
+
+			// Stripe Link is part of the card gateway, so we need to check separately if Link is enabled.
+			if ( Payment_Method::LINK === $payment_method && ! $this->is_payment_method_enabled( Payment_Method::LINK ) ) {
+				continue;
+			}
+
+			$types[] = $payment_method;
+		}
+
+		return $types;
+	}
+
+	/**
+	 * Checks if a payment method is enabled.
+	 *
+	 * @param string $payment_method The payment method to check.
+	 * @return bool True if the payment method is enabled, false otherwise.
+	 */
+	private function is_payment_method_enabled( $payment_method ) {
+		return in_array( $payment_method, WC_Payments::get_gateway()->get_upe_enabled_payment_method_ids(), true );
+	}
+
 	/**
 	 * Delete token from Stripe.
 	 *

tests/unit/test-class-wc-payments-token-service.php

@@ -489,17 +489,23 @@ public function test_woocommerce_get_customer_payment_tokens_multiple_tokens_mul
 		$gateway->settings['upe_enabled_payment_method_ids'] = $payment_methods;
 
 		// Array keys should match the database ID of the token.
-		$tokens = [
+		$card_tokens        = [
 			1 => $this->generate_card_token( 'pm_111', 1 ),
 			2 => $this->generate_card_token( 'pm_222', 2 ),
+		];
+		$sepa_tokens        = [
 			3 => $this->generate_sepa_token( 'pm_333', 3 ),
 			4 => $this->generate_sepa_token( 'pm_444', 4 ),
+		];
+		$stripe_link_tokens = [
 			5 => $this->generate_link_token( 'pm_555', 5 ),
 			6 => $this->generate_link_token( 'pm_666', 6 ),
 		];
 
+		$all_saved_tokens = $card_tokens + $sepa_tokens + $stripe_link_tokens;
+
 		$this->mock_customer_service
-			->expects( $this->once() )
+			->expects( $this->exactly( 2 ) )
 			->method( 'get_customer_id_by_user_id' )
 			->willReturn( $customer_id );
 
@@ -509,28 +515,34 @@ public function test_woocommerce_get_customer_payment_tokens_multiple_tokens_mul
 			->method( 'get_payment_methods_for_customer' )
 			->withConsecutive(
 				[ $customer_id, Payment_Method::CARD ],
-				[ $customer_id, Payment_Method::SEPA ],
 				[ $customer_id, Payment_Method::LINK ]
 			)
 			->willReturnOnConsecutiveCalls(
 				[
 					$this->generate_card_pm_response( 'pm_111' ),
 					$this->generate_card_pm_response( 'pm_222' ),
 				],
-				[
-					$this->generate_sepa_pm_response( 'pm_333' ),
-					$this->generate_sepa_pm_response( 'pm_444' ),
-				],
 				[
 					$this->generate_link_pm_response( 'pm_555' ),
 					$this->generate_link_pm_response( 'pm_666' ),
+				],
+				[
+					$this->generate_sepa_pm_response( 'pm_333' ),
+					$this->generate_sepa_pm_response( 'pm_444' ),
 				]
 			);
 
-		$result = $this->token_service->woocommerce_get_customer_payment_tokens( $tokens, 1, 'woocommerce_payments' );
+		$card_and_link_result = $this->token_service->woocommerce_get_customer_payment_tokens( $all_saved_tokens, 1, WC_Payment_Gateway_WCPay::GATEWAY_ID );
+		$sepa_result          = $this->token_service->woocommerce_get_customer_payment_tokens( $all_saved_tokens, 1, WC_Payment_Gateway_WCPay::GATEWAY_ID . '_' . Payment_Method::SEPA );
+
 		$this->assertSame(
-			array_keys( $tokens ),
-			array_keys( $result )
+			array_keys( $card_tokens + $stripe_link_tokens ),
+			array_keys( $card_and_link_result )
+		);
+
+		$this->assertSame(
+			array_keys( $sepa_tokens ),
+			array_keys( $sepa_result )
 		);
 	}
 
@@ -656,51 +668,6 @@ public function test_woocommerce_get_customer_payment_tokens_not_added_twice_for
 		$this->assertEquals( 'pm_444', $result_tokens[3]->get_token() );
 	}
 
-	public function test_woocommerce_get_customer_payment_tokens_not_added_from_different_gateway() {
-		$this->mock_cache->method( 'get' )->willReturn( [ 'is_deferred_intent_creation_upe_enabled' => true ] );
-		$gateway_id      = WC_Payment_Gateway_WCPay::GATEWAY_ID;
-		$tokens          = [];
-		$payment_methods = [ Payment_Method::CARD, Payment_Method::SEPA ];
-
-		$gateway = WC_Payments::get_gateway();
-		$gateway->settings['upe_enabled_payment_method_ids'] = $payment_methods;
-
-		$this->mock_customer_service
-			->expects( $this->any() )
-			->method( 'get_customer_id_by_user_id' )
-			->willReturn( 'cus_12345' );
-
-		$this->mock_customer_service
-			->expects( $this->exactly( 2 ) )
-			->method( 'get_payment_methods_for_customer' )
-			->withConsecutive(
-				[ 'cus_12345', Payment_Method::CARD ],
-				[ 'cus_12345', Payment_Method::SEPA ]
-			)
-			->willReturnOnConsecutiveCalls(
-				[
-					$this->generate_card_pm_response( 'pm_mock0' ),
-					$this->generate_card_pm_response( 'pm_222' ),
-				],
-				[
-					$this->generate_sepa_pm_response( 'other_gateway_pm_111' ),
-					$this->generate_sepa_pm_response( 'other_gateway_pm_222' ),
-					$this->generate_sepa_pm_response( 'other_gateway_pm_333' ),
-					$this->generate_sepa_pm_response( 'other_gateway_pm_444' ),
-					$this->generate_sepa_pm_response( 'other_gateway_pm_555' ),
-				]
-			);
-
-		$result        = $this->token_service->woocommerce_get_customer_payment_tokens( $tokens, 1, $gateway_id );
-		$result_tokens = array_values( $result );
-
-		$this->assertEquals( 2, count( $result_tokens ) );
-		$this->assertEquals( $gateway_id, $result_tokens[0]->get_gateway_id() );
-		$this->assertEquals( $gateway_id, $result_tokens[1]->get_gateway_id() );
-		$this->assertEquals( 'pm_mock0', $result_tokens[0]->get_token() );
-		$this->assertEquals( 'pm_222', $result_tokens[1]->get_token() );
-	}
-
 	public function test_woocommerce_get_customer_payment_tokens_payment_methods_only_for_retrievable_types() {
 		$enabled_upe_payment_methods = [
 			Payment_Method::CARD,
@@ -716,7 +683,6 @@ public function test_woocommerce_get_customer_payment_tokens_payment_methods_onl
 		$gateway                     = WC_Payments::get_gateway();
 		$gateway->settings['upe_enabled_payment_method_ids'] = $enabled_upe_payment_methods;
 		$tokens      = [];
-		$gateway_id  = 'woocommerce_payments';
 		$customer_id = 'cus_12345';
 
 		$this->mock_customer_service
@@ -729,22 +695,23 @@ public function test_woocommerce_get_customer_payment_tokens_payment_methods_onl
 			->method( 'get_payment_methods_for_customer' )
 			->withConsecutive(
 				[ $customer_id, Payment_Method::CARD ],
+				[ $customer_id, Payment_Method::LINK ],
 				[ $customer_id, Payment_Method::SEPA ],
-				[ $customer_id, Payment_Method::LINK ]
 			)
 			->willReturnOnConsecutiveCalls(
 				[
 					$this->generate_card_pm_response( 'pm_mock0' ),
 				],
 				[
-					$this->generate_sepa_pm_response( 'pm_mock_2' ),
+					$this->generate_link_pm_response( 'pm_mock_3' ),
 				],
 				[
-					$this->generate_link_pm_response( 'pm_mock_3' ),
-				]
+					$this->generate_sepa_pm_response( 'pm_mock_2' ),
+				],
 			);
 
-		$this->token_service->woocommerce_get_customer_payment_tokens( $tokens, 1, $gateway_id );
+		$this->token_service->woocommerce_get_customer_payment_tokens( $tokens, 1, WC_Payment_Gateway_WCPay::GATEWAY_ID );
+		$this->token_service->woocommerce_get_customer_payment_tokens( $tokens, 1, WC_Payment_Gateway_WCPay::GATEWAY_ID . '_' . Payment_Method::SEPA );
 	}
 
 	/**