[RSDK-11532] Add create_invoice_and_charge_immediately function (#971)

Author: edobranov

src/viam/app/billing_client.py

@@ -13,6 +13,8 @@
     GetInvoicesSummaryResponse,
     GetOrgBillingInformationRequest,
     GetOrgBillingInformationResponse,
+    CreateInvoiceAndChargeImmediatelyRequest,
+    CreateInvoiceAndChargeImmediatelyResponse,
 )
 
 LOGGER = logging.getLogger(__name__)
@@ -141,3 +143,20 @@ async def get_org_billing_information(self, org_id: str, timeout: Optional[float
         """
         request = GetOrgBillingInformationRequest(org_id=org_id)
         return await self._billing_client.GetOrgBillingInformation(request, metadata=self._metadata, timeout=timeout)
+
+    async def create_invoice_and_charge_immediately(self, org_id_to_charge: str, amount: float, description: Optional[str] = None, org_id_for_branding: Optional[str] = None) -> None:
+        """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.
+
+        ::
+
+            await billing_client.create_invoice_and_charge_immediately("<ORG-ID-TO-CHARGE>", <AMOUNT>, <DESCRIPTION>, "<ORG-ID-FOR-BRANDING>")
+
+        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
+        """
+        request = CreateInvoiceAndChargeImmediatelyRequest(org_id_to_charge=org_id_to_charge, amount=amount, description=description, org_id_for_branding=org_id_for_branding)
+        _: CreateInvoiceAndChargeImmediatelyResponse = await self._billing_client.CreateInvoiceAndChargeImmediately(request, metadata=self._metadata)

tests/mocks/services.py

@@ -191,6 +191,8 @@
     GetInvoicesSummaryResponse,
     GetOrgBillingInformationRequest,
     GetOrgBillingInformationResponse,
+    CreateInvoiceAndChargeImmediatelyRequest,
+    CreateInvoiceAndChargeImmediatelyResponse,
     UnimplementedBillingServiceBase,
 )
 from viam.proto.app.data import (
@@ -1297,6 +1299,14 @@ async def GetOrgBillingInformation(self, stream: Stream[GetOrgBillingInformation
         self.org_id = request.org_id
         await stream.send_message(self.billing_info)
 
+    async def CreateInvoiceAndChargeImmediately(self, stream: Stream[CreateInvoiceAndChargeImmediatelyRequest, CreateInvoiceAndChargeImmediatelyResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        self.org_id_to_charge = request.org_id_to_charge
+        self.amount = request.amount
+        self.description = request.description
+        self.org_id_for_branding = request.org_id_for_branding
+        await stream.send_message(CreateInvoiceAndChargeImmediatelyResponse())
 
 class MockApp(UnimplementedAppServiceBase):
     def __init__(

tests/test_billing_client.py

@@ -8,6 +8,7 @@
     GetInvoicesSummaryResponse,
     GetOrgBillingInformationResponse,
     InvoiceSummary,
+    CreateInvoiceAndChargeImmediatelyResponse,
 )
 
 from .mocks.services import MockBilling
@@ -105,3 +106,12 @@ async def test_get_org_billing_information(self, service: MockBilling):
             org_billing_info = await client.get_org_billing_information(org_id=org_id)
             assert org_billing_info == ORG_BILLING_INFO
             assert service.org_id == org_id
+
+    async def test_create_invoice_and_charge_immediately(self, service: MockBilling):
+        async with ChannelFor([service]) as channel:
+            client = BillingClient(channel, BILLING_SERVICE_METADATA)
+            _: CreateInvoiceAndChargeImmediatelyResponse = await client.create_invoice_and_charge_immediately(
+                org_id_to_charge="foo",
+                amount=42.42,
+                description="A short description.",
+                org_id_for_branding="bar")