RSDK-4907 - Add billing client (#471)

Author: stuqdog

src/viam/app/app_client.py

@@ -656,6 +656,7 @@ async def get_robot_part(self, robot_part_id: str, dest: Optional[str] = None, i
             try:
                 file = open(dest, "w")
                 file.write(f"{json.dumps(json.loads(response.config_json), indent=indent)}")
+                file.flush()
             except Exception as e:
                 LOGGER.error(f"Failed to write config JSON to file {dest}", exc_info=e)
 
@@ -716,6 +717,7 @@ async def get_robot_part_logs(
                     file_name = log.caller["File"] + ":" + str(int(log.caller["Line"]))
                     message = log.message
                     file.write(f"{time}\t{level}\t{logger_name}\t{file_name:<64}{message}\n")
+                    file.flush()
             except Exception as e:
                 LOGGER.error(f"Failed to write robot part from robot part with ID [{robot_part_id}]logs to file {dest}", exc_info=e)
 

src/viam/app/billing_client.py

@@ -0,0 +1,90 @@
+from typing import Mapping, Optional
+
+from grpclib.client import Channel
+
+from viam import logging
+from viam.proto.app.billing import (
+        BillingServiceStub,
+        GetCurrentMonthUsageRequest,
+        GetCurrentMonthUsageResponse,
+        GetInvoicePdfRequest,
+        GetInvoicePdfResponse,
+        GetInvoicesSummaryRequest,
+        GetInvoicesSummaryResponse,
+        GetOrgBillingInformationRequest,
+        GetOrgBillingInformationResponse,
+)
+
+LOGGER = logging.getLogger(__name__)
+
+
+class BillingClient:
+    """gRPC client for retrieving billing data from app.
+
+    Constructor is used by `ViamClient` to instantiate relevant service stubs. Calls to
+    `BillingClient` methods should be made through `ViamClient`.
+    """
+
+    def __init__(self, channel: Channel, metadata: Mapping[str, str]):
+        """Create a `BillingClient` that maintains a connection to app.
+
+        Args:
+            channel (grpclib.client.Channel): Connection to app.
+            metadata (Mapping[str, str]): Required authorization token to send requests to app.
+        """
+        self._metadata = metadata
+        self._billing_client = BillingServiceStub(channel)
+        self._channel = channel
+
+    _billing_client: BillingServiceStub
+    _channel: Channel
+    _metadata: Mapping[str, str]
+
+    async def get_current_month_usage(self, org_id: str, timeout: Optional[float] = None) -> GetCurrentMonthUsageResponse:
+        """Access data usage information for the current month for a given organization.
+
+        Args:
+            org_id (str): the ID of the organization to request usage data for
+
+        Returns:
+            viam.proto.app.billing.GetCurrentMonthUsageResponse: Current month usage information
+        """
+        request = GetCurrentMonthUsageRequest(org_id=org_id)
+        return await self._billing_client.GetCurrentMonthUsage(request, metadata=self._metadata, timeout=timeout)
+
+    async def get_invoice_pdf(self, invoice_id: str, org_id: str, dest: str, timeout: Optional[float] = None) -> None:
+        """Access invoice PDF data and optionally save it to a provided file path.
+
+        Args:
+            invoice_id (str): the ID of the invoice being requested
+            org_id (str): the ID of the org to request data from
+            dest (str): filepath to save the invoice to
+        """
+        request = GetInvoicePdfRequest(id=invoice_id, org_id=org_id)
+        response: GetInvoicePdfResponse = await self._billing_client.GetInvoicePdf(request, metadata=self._metadata, timeout=timeout)
+        data: bytes = response[0].chunk
+        with open(dest, "wb") as file:
+            file.write(data)
+
+    async def get_invoices_summary(self, org_id: str, timeout: Optional[float] = None) -> GetInvoicesSummaryResponse:
+        """Access total outstanding balance plus invoice summaries for a given org.
+
+        Args:
+            org_id (str): the ID of the org to request data for
+
+        Returns:
+            viam.proto.app.billing.GetInvoicesSummaryResponse: Summary of org invoices
+        """
+        request = GetInvoicesSummaryRequest(org_id=org_id)
+        return await self._billing_client.GetInvoicesSummary(request, metadata=self._metadata, timeout=timeout)
+
+    async def get_org_billing_information(self, org_id: str, timeout: Optional[float] = None) -> GetOrgBillingInformationResponse:
+        """Access billing information (payment method, billing tier, etc.) for a given org.
+
+        Args:
+            org_id (str): the ID of the org to request data for
+
+        Returns:
+            viam.proto.app.billing.GetOrgBillingInformationResponse: The org billing information"""
+        request = GetOrgBillingInformationRequest(org_id=org_id)
+        return await self._billing_client.GetOrgBillingInformation(request, metadata=self._metadata, timeout=timeout)

src/viam/app/data_client.py

@@ -170,6 +170,7 @@ async def tabular_data_by_filter(
             try:
                 file = open(dest, "w")
                 file.write(f"{[str(d) for d in data]}")
+                file.flush()
             except Exception as e:
                 LOGGER.error(f"Failed to write tabular data to file {dest}", exc_info=e)
         return data
@@ -222,6 +223,7 @@ async def binary_data_by_filter(
             try:
                 file = open(dest, "w")
                 file.write(f"{[str(d) for d in data]}")
+                file.flush()
             except Exception as e:
                 LOGGER.error(f"Failed to write binary data to file {dest}", exc_info=e)
 
@@ -256,6 +258,7 @@ async def binary_data_by_ids(
             try:
                 file = open(dest, "w")
                 file.write(f"{response.data}")
+                file.flush()
             except Exception as e:
                 LOGGER.error(f"Failed to write binary data to file {dest}", exc_info=e)
         return [DataClient.BinaryData(data.binary, data.metadata) for data in response.data]

src/viam/app/viam_client.py

@@ -5,6 +5,7 @@
 
 from viam import logging
 from viam.app.app_client import AppClient
+from viam.app.billing_client import BillingClient
 from viam.app.data_client import DataClient
 from viam.app.ml_training_client import MLTrainingClient
 from viam.rpc.dial import DialOptions, _dial_app, _get_access_token
@@ -74,6 +75,11 @@ def ml_training_client(self) -> MLTrainingClient:
         """Instantiate and return a `MLTrainingClient` used to make `ml_training` method calls."""
         return MLTrainingClient(self._channel, self._metadata)
 
+    @property
+    def billing_client(self) -> BillingClient:
+        """Instantiate and return a `BillingClient` used to make `billing` method calls."""
+        return BillingClient(self._channel, self._metadata)
+
     def close(self):
         """Close opened channels used for the various service stubs initialized."""
         if self._closed:

tests/mocks/services.py

@@ -211,6 +211,27 @@
     SubmitTrainingJobResponse,
     TrainingJobMetadata,
 )
+from viam.proto.app.billing import (
+    BillingServiceBase,
+    GetCurrentMonthUsageRequest,
+    GetCurrentMonthUsageResponse,
+    GetInvoicePdfRequest,
+    GetInvoicePdfResponse,
+    GetInvoicesSummaryRequest,
+    GetInvoicesSummaryResponse,
+    GetOrgBillingInformationRequest,
+    GetOrgBillingInformationResponse,
+    GetBillingSummaryRequest,
+    GetBillingSummaryResponse,
+    GetCurrentMonthUsageSummaryRequest,
+    GetCurrentMonthUsageSummaryResponse,
+    GetInvoiceHistoryRequest,
+    GetInvoiceHistoryResponse,
+    GetItemizedInvoiceRequest,
+    GetItemizedInvoiceResponse,
+    GetUnpaidBalanceRequest,
+    GetUnpaidBalanceResponse,
+)
 from viam.proto.common import DoCommandRequest, DoCommandResponse, GeoObstacle, GeoPoint, PointCloudObject, Pose, PoseInFrame, ResourceName
 from viam.proto.service.mlmodel import (
     FlatTensor,
@@ -818,6 +839,64 @@ async def CancelTrainingJob(self, stream: Stream[CancelTrainingJobRequest, Cance
         await stream.send_message(CancelTrainingJobResponse())
 
 
+class MockBilling(BillingServiceBase):
+    def __init__(
+        self,
+        pdf: bytes,
+        curr_month_usage: GetCurrentMonthUsageResponse,
+        invoices_summary: GetInvoicesSummaryResponse,
+        billing_info: GetOrgBillingInformationResponse,
+    ):
+        self.pdf = pdf
+        self.curr_month_usage = curr_month_usage
+        self.invoices_summary = invoices_summary
+        self.billing_info = billing_info
+
+    async def GetCurrentMonthUsage(self, stream: Stream[GetCurrentMonthUsageRequest, GetCurrentMonthUsageResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        self.org_id = request.org_id
+        await stream.send_message(self.curr_month_usage)
+
+    async def GetInvoicePdf(self, stream: Stream[GetInvoicePdfRequest, GetInvoicePdfResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        self.org_id = request.org_id
+        self.invoice_id = request.id
+        response = GetInvoicePdfResponse(chunk=self.pdf)
+        await stream.send_message(response)
+
+    async def GetInvoicesSummary(self, stream: Stream[GetInvoicesSummaryRequest, GetInvoicePdfResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        self.org_id = request.org_id
+        await stream.send_message(self.invoices_summary)
+
+    async def GetOrgBillingInformation(self, stream: Stream[GetOrgBillingInformationRequest, GetOrgBillingInformationResponse]) -> None:
+        request = await stream.recv_message()
+        assert request is not None
+        self.org_id = request.org_id
+        await stream.send_message(self.billing_info)
+
+    async def GetBillingSummary(self, stream: Stream[GetBillingSummaryRequest, GetBillingSummaryResponse]) -> None:
+        raise NotImplementedError()
+
+    async def GetCurrentMonthUsageSummary(
+        self,
+        stream: Stream[GetCurrentMonthUsageSummaryRequest, GetCurrentMonthUsageSummaryResponse],
+    ) -> None:
+        raise NotImplementedError()
+
+    async def GetInvoiceHistory(self, stream: Stream[GetInvoiceHistoryRequest, GetInvoiceHistoryResponse]) -> None:
+        raise NotImplementedError()
+
+    async def GetItemizedInvoice(self, stream: Stream[GetItemizedInvoiceRequest, GetItemizedInvoiceResponse]) -> None:
+        raise NotImplementedError()
+
+    async def GetUnpaidBalance(self, stream: Stream[GetUnpaidBalanceRequest, GetUnpaidBalanceResponse]) -> None:
+        raise NotImplementedError()
+
+
 class MockApp(AppServiceBase):
     def __init__(
         self,

tests/test_billing_client.py

@@ -0,0 +1,112 @@
+import pytest
+
+from google.protobuf.timestamp_pb2 import Timestamp
+from grpclib.testing import ChannelFor
+
+from viam.app.billing_client import BillingClient
+from viam.proto.app.billing import (
+    GetCurrentMonthUsageResponse,
+    GetInvoicesSummaryResponse,
+    GetOrgBillingInformationResponse,
+    InvoiceSummary,
+)
+
+from .mocks.services import MockBilling
+
+PDF = b'abc123'
+CLOUD_STORAGE_USAGE_COST = 100.0
+DATA_UPLOAD_USAGE_COST = 101.0
+DATA_EGRES_USAGE_COST = 102.0
+REMOTE_CONTROL_USAGE_COST = 103.0
+STANDARD_COMPUTE_USAGE_COST = 104.0
+DISCOUNT_AMOUNT = 0.0
+TOTAL_USAGE_WITH_DISCOUNT = 105.0
+TOTAL_USAGE_WITHOUT_DISCOUNT = 106.0
+OUTSTANDING_BALANCE = 1000.0
+SECONDS_START = 978310861
+NANOS_START = 0
+SECONDS_END = 998310861
+NANOS_END = 0
+SECONDS_PAID = 988310861
+NANOS_PAID = 0
+START_TS = Timestamp(seconds=SECONDS_START, nanos=NANOS_END)
+PAID_DATE_TS = Timestamp(seconds=SECONDS_PAID, nanos=NANOS_PAID)
+END_TS = Timestamp(seconds=SECONDS_END, nanos=NANOS_END)
+INVOICE_ID = "invoice"
+STATUS = "status"
+PAYMENT_TYPE = 1
+EMAIL = "[email protected]"
+BILLING_TIER = "tier"
+INVOICE = InvoiceSummary(
+    id=INVOICE_ID,
+    invoice_date=START_TS,
+    invoice_amount=OUTSTANDING_BALANCE,
+    status=STATUS,
+    due_date=END_TS,
+    paid_date=PAID_DATE_TS,
+)
+INVOICES = [INVOICE]
+CURR_MONTH_USAGE = GetCurrentMonthUsageResponse(
+    start_date=START_TS,
+    end_date=END_TS,
+    cloud_storage_usage_cost=CLOUD_STORAGE_USAGE_COST,
+    data_upload_usage_cost=DATA_UPLOAD_USAGE_COST,
+    data_egres_usage_cost=DATA_EGRES_USAGE_COST,
+    remote_control_usage_cost=REMOTE_CONTROL_USAGE_COST,
+    standard_compute_usage_cost=STANDARD_COMPUTE_USAGE_COST,
+    discount_amount=DISCOUNT_AMOUNT,
+    total_usage_with_discount=TOTAL_USAGE_WITH_DISCOUNT,
+    total_usage_without_discount=TOTAL_USAGE_WITHOUT_DISCOUNT,
+)
+INVOICES_SUMMARY = GetInvoicesSummaryResponse(outstanding_balance=OUTSTANDING_BALANCE, invoices=INVOICES)
+ORG_BILLING_INFO = GetOrgBillingInformationResponse(
+    type=PAYMENT_TYPE,
+    billing_email=EMAIL,
+    billing_tier=BILLING_TIER,
+)
+
+AUTH_TOKEN = "auth_token"
+BILLING_SERVICE_METADATA = {"authorization": f"Bearer {AUTH_TOKEN}"}
+
+
+@pytest.fixture(scope="function")
+def service() -> MockBilling:
+    return MockBilling(
+        pdf=PDF,
+        curr_month_usage=CURR_MONTH_USAGE,
+        invoices_summary=INVOICES_SUMMARY,
+        billing_info=ORG_BILLING_INFO,
+    )
+
+
+class TestClient:
+    @pytest.mark.asyncio
+    async def test_get_current_month_usage(self, service: MockBilling):
+        async with ChannelFor([service]) as channel:
+            org_id = "foo"
+            client = BillingClient(channel, BILLING_SERVICE_METADATA)
+            curr_month_usage = await client.get_current_month_usage(org_id=org_id)
+            assert curr_month_usage == CURR_MONTH_USAGE
+            assert service.org_id == org_id
+
+    @pytest.mark.asyncio
+    async def test_get_invoice_pdf(self, service: MockBilling):
+        assert True
+
+    @pytest.mark.asyncio
+    async def test_get_invoices_summary(self, service: MockBilling):
+        async with ChannelFor([service]) as channel:
+            org_id = "bar"
+            client = BillingClient(channel, BILLING_SERVICE_METADATA)
+            invoices_summary = await client.get_invoices_summary(org_id=org_id)
+            assert invoices_summary == INVOICES_SUMMARY
+            assert service.org_id == org_id
+
+    @pytest.mark.asyncio
+    async def test_get_org_billing_information(self, service: MockBilling):
+        async with ChannelFor([service]) as channel:
+            org_id = "baz"
+            client = BillingClient(channel, BILLING_SERVICE_METADATA)
+            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