AI update based on proto changes (#1016)

Co-authored-by: njooma <[email protected]>
Co-authored-by: Naveed Jooma <[email protected]>

Author: viambot

src/viam/app/billing_client.py

@@ -143,7 +143,12 @@ async def get_org_billing_information(self, org_id: str, timeout: Optional[float
         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
+        self,
+        org_id_to_charge: str,
+        amount: float,
+        description: Optional[str] = None,
+        org_id_for_branding: Optional[str] = None,
+        disable_email: bool = False,
     ) -> 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.
@@ -157,9 +162,14 @@ async def create_invoice_and_charge_immediately(
             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
         """
         request = CreateInvoiceAndChargeImmediatelyRequest(
-            org_id_to_charge=org_id_to_charge, amount=amount, description=description, org_id_for_branding=org_id_for_branding
+            org_id_to_charge=org_id_to_charge,
+            amount=amount,
+            description=description,
+            org_id_for_branding=org_id_for_branding,
+            disable_email=disable_email,
         )
         _: CreateInvoiceAndChargeImmediatelyResponse = await self._billing_client.CreateInvoiceAndChargeImmediately(
             request, metadata=self._metadata

tests/mocks/services.py

@@ -1306,6 +1306,7 @@ def __init__(
         self.curr_month_usage = curr_month_usage
         self.invoices_summary = invoices_summary
         self.billing_info = billing_info
+        self.disable_email: bool = False
 
     async def GetCurrentMonthUsage(self, stream: Stream[GetCurrentMonthUsageRequest, GetCurrentMonthUsageResponse]) -> None:
         request = await stream.recv_message()
@@ -1342,6 +1343,7 @@ async def CreateInvoiceAndChargeImmediately(
         self.amount = request.amount
         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())
 
 
@@ -1937,11 +1939,7 @@ async def list_uuids(
         return [b"uuid1", b"uuid2", b"uuid3"]
 
     async def get_transform(
-        self,
-        uuid: bytes,
-        *,
-        extra: Optional[Mapping[str, Any]] = None,
-        timeout: Optional[float] = None,
+        self, uuid: bytes, *, extra: Optional[Mapping[str, Any]] = None, timeout: Optional[float] = None
     ) -> "Transform":
         self.uuid = uuid
         self.extra = extra
@@ -1955,10 +1953,7 @@ async def get_transform(
         )
 
     async def stream_transform_changes(
-        self,
-        *,
-        extra: Optional[Mapping[str, Any]] = None,
-        timeout: Optional[float] = None,
+        self, *, extra: Optional[Mapping[str, Any]] = None, timeout: Optional[float] = None
     ) -> AsyncGenerator[StreamTransformChangesResponse, None]:
         self.extra = extra
         self.timeout = timeout

tests/test_billing_client.py

@@ -113,13 +113,16 @@ async def test_create_invoice_and_charge_immediately(self, service: MockBilling)
             org_id = "foo"
             description = "A short description"
             org_id_for_branding = "bar"
+            disable_email = True
             _: CreateInvoiceAndChargeImmediatelyResponse = 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 service.org_id_to_charge == org_id
             assert service.description == description
             assert service.org_id_for_branding == org_id_for_branding
             assert service.amount == OUTSTANDING_BALANCE
+            assert service.disable_email == disable_email