update logic to support multiple subscriptions

Author: ajay-sentry

billing/tests/test_views.py

@@ -1643,6 +1643,119 @@ class MockPaymentMethod:
             "sub_123", default_payment_method=payment_method_retrieve_mock.return_value
         )
 
+    @patch("services.billing.stripe.PaymentMethod.attach")
+    @patch("services.billing.stripe.Customer.modify")
+    @patch("services.billing.stripe.Subscription.modify")
+    @patch("services.billing.stripe.PaymentMethod.retrieve")
+    def test_check_and_handle_delayed_notification_payment_methods_no_subscription(
+        self,
+        payment_method_retrieve_mock,
+        subscription_modify_mock,
+        customer_modify_mock,
+        payment_method_attach_mock,
+    ):
+        class MockPaymentMethod:
+            type = "us_bank_account"
+            us_bank_account = {}
+            id = "pm_123"
+
+        payment_method_retrieve_mock.return_value = MockPaymentMethod()
+
+        self.owner.stripe_subscription_id = None
+        self.owner.stripe_customer_id = "cus_123"
+        self.owner.save()
+
+        handler = StripeWebhookHandler()
+        handler._check_and_handle_delayed_notification_payment_methods(
+            "cus_123", "pm_123"
+        )
+
+        payment_method_retrieve_mock.assert_called_once_with("pm_123")
+        payment_method_attach_mock.assert_not_called()
+        customer_modify_mock.assert_not_called()
+        subscription_modify_mock.assert_not_called()
+
+    @patch("services.billing.stripe.PaymentMethod.attach")
+    @patch("services.billing.stripe.Customer.modify")
+    @patch("services.billing.stripe.Subscription.modify")
+    @patch("services.billing.stripe.PaymentMethod.retrieve")
+    def test_check_and_handle_delayed_notification_payment_methods_no_customer(
+        self,
+        payment_method_retrieve_mock,
+        subscription_modify_mock,
+        customer_modify_mock,
+        payment_method_attach_mock,
+    ):
+        class MockPaymentMethod:
+            type = "us_bank_account"
+            us_bank_account = {}
+            id = "pm_123"
+
+        payment_method_retrieve_mock.return_value = MockPaymentMethod()
+
+        handler = StripeWebhookHandler()
+        handler._check_and_handle_delayed_notification_payment_methods(
+            "cus_123", "pm_123"
+        )
+
+        payment_method_retrieve_mock.assert_called_once_with("pm_123")
+        payment_method_attach_mock.assert_not_called()
+        customer_modify_mock.assert_not_called()
+        subscription_modify_mock.assert_not_called()
+
+    @patch("logging.Logger.error")
+    @patch("services.billing.stripe.PaymentMethod.attach")
+    @patch("services.billing.stripe.Customer.modify")
+    @patch("services.billing.stripe.Subscription.modify")
+    @patch("services.billing.stripe.PaymentMethod.retrieve")
+    def test_check_and_handle_delayed_notification_payment_methods_multiple_subscriptions(
+        self,
+        payment_method_retrieve_mock,
+        subscription_modify_mock,
+        customer_modify_mock,
+        payment_method_attach_mock,
+    ):
+        class MockPaymentMethod:
+            type = "us_bank_account"
+            us_bank_account = {}
+            id = "pm_123"
+
+        payment_method_retrieve_mock.return_value = MockPaymentMethod()
+
+        self.owner.stripe_subscription_id = "sub_123"
+        self.owner.stripe_customer_id = "cus_123"
+        self.owner.save()
+
+        OwnerFactory(stripe_subscription_id="sub_124", stripe_customer_id="cus_123")
+
+        handler = StripeWebhookHandler()
+        handler._check_and_handle_delayed_notification_payment_methods(
+            "cus_123", "pm_123"
+        )
+
+        payment_method_retrieve_mock.assert_called_once_with("pm_123")
+        payment_method_attach_mock.assert_called_once_with(
+            payment_method_retrieve_mock.return_value, customer="cus_123"
+        )
+        customer_modify_mock.assert_called_once_with(
+            "cus_123",
+            invoice_settings={
+                "default_payment_method": payment_method_retrieve_mock.return_value
+            },
+        )
+        subscription_modify_mock.assert_has_calls(
+            [
+                call(
+                    "sub_123",
+                    default_payment_method=payment_method_retrieve_mock.return_value,
+                ),
+                call(
+                    "sub_124",
+                    default_payment_method=payment_method_retrieve_mock.return_value,
+                ),
+            ]
+        )
+
     @patch(
         "billing.views.StripeWebhookHandler._check_and_handle_delayed_notification_payment_methods"
     )

billing/views.py

@@ -539,7 +539,6 @@ def _check_and_handle_delayed_notification_payment_methods(
         When verification succeeds, this attaches the payment method to the customer and sets
         it as the default payment method for both the customer and subscription.
         """
-        owner = Owner.objects.get(stripe_customer_id=customer_id)
         payment_method = stripe.PaymentMethod.retrieve(payment_method_id)
 
         is_us_bank_account = payment_method.type == "us_bank_account" and hasattr(
@@ -548,19 +547,34 @@ def _check_and_handle_delayed_notification_payment_methods(
 
         should_set_as_default = is_us_bank_account
 
+        # attach the payment method + set as default on the invoice and subscription
         if should_set_as_default:
-            # attach the payment method + set as default on the invoice and subscription
-            stripe.PaymentMethod.attach(
-                payment_method, customer=owner.stripe_customer_id
-            )
-            stripe.Customer.modify(
-                owner.stripe_customer_id,
-                invoice_settings={"default_payment_method": payment_method},
-            )
-            stripe.Subscription.modify(
-                owner.stripe_subscription_id, default_payment_method=payment_method
+            # retrieve the number of owners to update
+            owners = Owner.objects.filter(
+                stripe_customer_id=customer_id, stripe_subscription_id__isnull=False
             )
 
+            if owners.exists():
+                # Even if multiple results are returned, these two stripe calls are
+                # just for a single customer
+                stripe.PaymentMethod.attach(payment_method, customer=customer_id)
+                stripe.Customer.modify(
+                    customer_id,
+                    invoice_settings={"default_payment_method": payment_method},
+                )
+
+                # But this one is for each subscription an owner may have
+                for owner in owners:
+                    stripe.Subscription.modify(
+                        owner.stripe_subscription_id,
+                        default_payment_method=payment_method,
+                    )
+            else:
+                log.error(
+                    "No owners found with that customer_id, something went wrong",
+                    extra=dict(customer_id=customer_id),
+                )
+
     def payment_intent_succeeded(self, payment_intent: stripe.PaymentIntent) -> None:
         """
         Stripe payment_intent.succeeded webhook event is emitted when a