Return invoice ID for create_invoice_and_charge_immediately

Author: edobranov

src/viam/app/billing_client.py

@@ -149,20 +149,31 @@ async def create_invoice_and_charge_immediately(
         description: Optional[str] = None,
         org_id_for_branding: Optional[str] = None,
         disable_email: bool = False,
-    ) -> None:
+    ) -> CreateInvoiceAndChargeImmediatelyResponse:
         """Create a flat fee invoice and charge the organization on the spot. The caller must be an owner of the organization being charged.
-        This function blocks until payment is confirmed, but will time out after 2 minutes if there is no confirmation.
+        This function returns the invoice id once the payment intent is successfully sent for processing. Callers may poll the invoice for
+        its status using the `get_invoices_summary` function and the returned invoice id. The status will be "payment_processing" if the
+        payment is being processed, "paid" if it succeeds, or "outstanding" if it fails.
 
         ::
 
-            await billing_client.create_invoice_and_charge_immediately("<ORG-ID-TO-CHARGE>", <AMOUNT>, <DESCRIPTION>, "<ORG-ID-FOR-BRANDING>", False)
+            invoice_id = await billing_client.create_invoice_and_charge_immediately(
+                "<ORG-ID-TO-CHARGE>",
+                <AMOUNT>,
+                <DESCRIPTION>,
+                "<ORG-ID-FOR-BRANDING>",
+                False,
+            )
 
         Args:
             org_id_to_charge (str): the organization to charge
             amount (float): the amount to charge in dollars
             description (str): a short description of the charge to display on the invoice PDF (must be 100 characters or less)
             org_id_for_branding (str): the organization whose branding to use in the invoice confirmation email
             disable_email (bool): whether or not to disable sending an email confirmation for the invoice
+
+        Returns:
+            viam.proto.app.billing.CreateInvoiceAndChargeImmediatelyResponse: the invoice id
         """
         request = CreateInvoiceAndChargeImmediatelyRequest(
             org_id_to_charge=org_id_to_charge,
@@ -171,6 +182,4 @@ async def create_invoice_and_charge_immediately(
             org_id_for_branding=org_id_for_branding,
             disable_email=disable_email,
         )
-        _: CreateInvoiceAndChargeImmediatelyResponse = await self._billing_client.CreateInvoiceAndChargeImmediately(
-            request, metadata=self._metadata
-        )
+        return await self._billing_client.CreateInvoiceAndChargeImmediately(request, metadata=self._metadata)

tests/mocks/services.py

@@ -1301,11 +1301,13 @@ def __init__(
         curr_month_usage: GetCurrentMonthUsageResponse,
         invoices_summary: GetInvoicesSummaryResponse,
         billing_info: GetOrgBillingInformationResponse,
+        invoice_id_response: CreateInvoiceAndChargeImmediatelyResponse,
     ):
         self.pdf = pdf
         self.curr_month_usage = curr_month_usage
         self.invoices_summary = invoices_summary
         self.billing_info = billing_info
+        self.invoice_id_response = invoice_id_response
         self.disable_email: bool = False
 
     async def GetCurrentMonthUsage(self, stream: Stream[GetCurrentMonthUsageRequest, GetCurrentMonthUsageResponse]) -> None:
@@ -1344,7 +1346,7 @@ async def CreateInvoiceAndChargeImmediately(
         self.description = request.description
         self.org_id_for_branding = request.org_id_for_branding
         self.disable_email = request.disable_email
-        await stream.send_message(CreateInvoiceAndChargeImmediatelyResponse())
+        await stream.send_message(self.invoice_id_response)
 
 
 class MockApp(UnimplementedAppServiceBase):

tests/test_billing_client.py

@@ -64,6 +64,7 @@
     billing_email=EMAIL,
     billing_tier=BILLING_TIER,
 )
+INVOICE_ID_RESPONSE = CreateInvoiceAndChargeImmediatelyResponse(invoice_id=INVOICE_ID)
 
 AUTH_TOKEN = "auth_token"
 BILLING_SERVICE_METADATA = {"authorization": f"Bearer {AUTH_TOKEN}"}
@@ -76,6 +77,7 @@ def service() -> MockBilling:
         curr_month_usage=CURR_MONTH_USAGE,
         invoices_summary=INVOICES_SUMMARY,
         billing_info=ORG_BILLING_INFO,
+        invoice_id_response=INVOICE_ID_RESPONSE,
     )
 
 
@@ -114,13 +116,14 @@ async def test_create_invoice_and_charge_immediately(self, service: MockBilling)
             description = "A short description"
             org_id_for_branding = "bar"
             disable_email = True
-            _: CreateInvoiceAndChargeImmediatelyResponse = await client.create_invoice_and_charge_immediately(
+            invoice_id_response = await client.create_invoice_and_charge_immediately(
                 org_id_to_charge=org_id,
                 amount=OUTSTANDING_BALANCE,
                 description=description,
                 org_id_for_branding=org_id_for_branding,
                 disable_email=disable_email,
             )
+            assert invoice_id_response == INVOICE_ID_RESPONSE
             assert service.org_id_to_charge == org_id
             assert service.description == description
             assert service.org_id_for_branding == org_id_for_branding