Fix the failing test for UPE gateway

includes/class-wc-payments-upe-checkout.php

@@ -281,7 +281,7 @@ function() use ( $payment_fields ) {
 	 */
 	public function set_gateway( $payment_method_id ) {
 		if ( null !== $payment_method_id ) {
-			$this->gateway = WC_Payments::get_payment_gateway_by_id( $payment_method_id );
+			$this->gateway = $this->gateway->wc_payments_get_payment_gateway_by_id( $payment_method_id );
 		}
 	}
 

includes/payment-methods/class-upe-payment-gateway.php

@@ -1123,6 +1123,16 @@ public function wc_payments_get_payment_method_by_id( $payment_method_id ) {
 		return WC_Payments::get_payment_method_by_id( $payment_method_id );
 	}
 
+	/**
+	 * This function wraps WC_Payments::get_payment_gateway_by_id, useful for unit testing.
+	 *
+	 * @param string $payment_method_id Stripe payment method type ID.
+	 * @return false|UPE_Payment_Gateway Matching UPE Payment Gateway instance.
+	 */
+	public function wc_payments_get_payment_gateway_by_id( $payment_method_id ) {
+		return WC_Payments::get_payment_gateway_by_id( $payment_method_id );
+	}
+
 	/**
 	 * Returns the UPE payment method for the gateway.
 	 *

tests/unit/payment-methods/test-class-upe-payment-gateway.php

@@ -269,6 +269,7 @@ public function set_up() {
 						'get_upe_enabled_payment_method_statuses',
 						'is_payment_recurring',
 						'get_payment_method_ids_enabled_at_checkout',
+						'wc_payments_get_payment_gateway_by_id',
 					]
 				)
 				->getMock();
@@ -297,18 +298,22 @@ public function set_up() {
 	 * @return void
 	 */
 	public function test_payment_fields_outputs_fields() {
-		new WC_Payments_UPE_Checkout(
-			$this->mock_payment_gateways['card'],
-			$this->mock_platform_checkout_utilities,
-			$this->mock_wcpay_account,
-			$this->mock_customer_service,
-			$this->mock_payment_gateways
-		);
 		foreach ( $this->mock_payment_gateways as $payment_method_id => $mock_payment_gateway ) {
+			new WC_Payments_UPE_Checkout(
+				$mock_payment_gateway,
+				$this->mock_platform_checkout_utilities,
+				$this->mock_wcpay_account,
+				$this->mock_customer_service
+			);
+
 			$mock_payment_gateway
 				->method( 'get_payment_method_ids_enabled_at_checkout' )
 				->willReturn( [] );
 
+			$mock_payment_gateway
+				->method( 'wc_payments_get_payment_gateway_by_id' )
+				->willReturn( $mock_payment_gateway );
+
 			/**
 			* This tests each payment method output separately without concatenating the output
 			* into 1 single buffer. Each iteration has 1 assertion.
@@ -1538,8 +1543,7 @@ public function test_remove_link_payment_method_if_card_disabled() {
 			$mock_upe_gateway,
 			$this->mock_platform_checkout_utilities,
 			$this->mock_wcpay_account,
-			$this->mock_customer_service,
-			[]
+			$this->mock_customer_service
 		);
 
 		$this->assertSame( $upe_checkout->get_payment_fields_js_config()['paymentMethodsConfig'], [] );
@@ -1599,8 +1603,7 @@ public function test_link_payment_method_if_card_enabled() {
 			$mock_upe_gateway,
 			$this->mock_platform_checkout_utilities,
 			$this->mock_wcpay_account,
-			$this->mock_customer_service,
-			[]
+			$this->mock_customer_service
 		);
 
 		$this->assertSame(