RSDK-7215: add billing wrapper (#276)

src/app/billing-client.test.ts

@@ -0,0 +1,155 @@
+import { FakeTransportBuilder } from '@improbable-eng/grpc-web-fake-transport';
+import { Timestamp } from 'google-protobuf/google/protobuf/timestamp_pb';
+import { afterEach, describe } from 'vitest';
+import { beforeEach, expect, it, vi } from 'vitest';
+import { type ResponseStream } from '../gen/robot/v1/robot_pb_service';
+import { EventDispatcher } from '../events';
+import {
+  GetCurrentMonthUsageRequest,
+  GetCurrentMonthUsageResponse,
+  GetInvoicePdfResponse,
+  GetInvoicesSummaryRequest,
+  GetOrgBillingInformationRequest,
+  PaymentMethodType,
+} from '../gen/app/v1/billing_pb';
+import { BillingServiceClient } from '../gen/app/v1/billing_pb_service';
+import { BillingClient } from './billing-client';
+
+const SECONDS = 1;
+const NANOS = 2_000_000;
+const testStartDate = new Timestamp();
+testStartDate.setSeconds(SECONDS);
+testStartDate.setNanos(NANOS);
+const testEndDate = new Timestamp();
+testEndDate.setSeconds(SECONDS * 2);
+testEndDate.setNanos(NANOS);
+const testMonthUsage = {
+  cloudStorageUsageCost: 1,
+  dataUploadUsageCost: 2,
+  dataEgresUsageCost: 3,
+  remoteControlUsageCost: 4,
+  standardComputeUsageCost: 5,
+  discountAmount: 6,
+  totalUsageWithDiscount: 7,
+  totalUsageWithoutDiscount: 8,
+  startDate: testStartDate.toObject(),
+  endDate: testEndDate.toObject(),
+  start: new Date(SECONDS * 1000 + NANOS / 1_000_000),
+  end: new Date(SECONDS * 2000 + NANOS / 1_000_000),
+};
+const testInvoiceSummary = {
+  id: 'id',
+  invoiceAmount: 1,
+  status: 'status',
+};
+const testBillingInfo = {
+  type: PaymentMethodType.PAYMENT_METHOD_TYPE_UNSPECIFIED,
+  billingEmail: '[email protected]',
+  billingTier: 'platinum',
+};
+
+class TestResponseStream<T> extends EventDispatcher {
+  private stream: ResponseStream<any>;
+
+  constructor(stream: ResponseStream<any>) {
+    super();
+    this.stream = stream;
+  }
+
+  override on(
+    type: string,
+    handler: (message: any) => void
+  ): ResponseStream<T> {
+    super.on(type, handler);
+    return this;
+  }
+
+  cancel(): void {
+    this.listeners = {};
+    this.stream.cancel();
+  }
+}
+let getGetInvoicePdfStream: ResponseStream<GetInvoicePdfResponse>;
+let testGetInvoicePdfStream:
+  | TestResponseStream<GetInvoicePdfResponse>
+  | undefined;
+
+const subject = () =>
+  new BillingClient('fakeServiceHost', {
+    transport: new FakeTransportBuilder().build(),
+  });
+
+describe('BillingClient tests', () => {
+  beforeEach(() => {
+    testGetInvoicePdfStream = new TestResponseStream(getGetInvoicePdfStream);
+    vi.spyOn(BillingServiceClient.prototype, 'getCurrentMonthUsage')
+      // @ts-expect-error compiler is matching incorrect function signature
+      .mockImplementation((_req: GetCurrentMonthUsageRequest, _md, cb) => {
+        const response = new GetCurrentMonthUsageResponse();
+        response.setCloudStorageUsageCost(1);
+        response.setDataUploadUsageCost(2);
+        response.setDataEgresUsageCost(3);
+        response.setRemoteControlUsageCost(4);
+        response.setStandardComputeUsageCost(5);
+        response.setDiscountAmount(6);
+        response.setTotalUsageWithDiscount(7);
+        response.setTotalUsageWithoutDiscount(8);
+        response.setStartDate(testStartDate);
+        response.setEndDate(testEndDate);
+        cb(null, response);
+      });
+
+    vi.spyOn(BillingServiceClient.prototype, 'getOrgBillingInformation')
+      // @ts-expect-error compiler is matching incorrect function signature
+      .mockImplementation((_req: GetOrgBillingInformationRequest, _md, cb) => {
+        cb(null, { toObject: () => testBillingInfo });
+      });
+
+    vi.spyOn(BillingServiceClient.prototype, 'getInvoicesSummary')
+      // @ts-expect-error compiler is matching incorrect function signature
+      .mockImplementation((_req: GetInvoicesSummaryRequest, _md, cb) => {
+        cb(null, { toObject: () => testInvoiceSummary });
+      });
+
+    BillingServiceClient.prototype.getInvoicePdf = vi
+      .fn()
+      .mockImplementation(() => testGetInvoicePdfStream);
+  });
+
+  afterEach(() => {
+    testGetInvoicePdfStream = undefined;
+  });
+
+  it('getCurrentMonthUsage', async () => {
+    const response = await subject().getCurrentMonthUsage('orgId');
+    expect(response).toEqual(testMonthUsage);
+  });
+
+  it('getOrgBillingInformation', async () => {
+    const response = await subject().getOrgBillingInformation('orgId');
+    expect(response).toEqual(testBillingInfo);
+  });
+
+  it('getInvoicesSummary', async () => {
+    const response = await subject().getInvoicesSummary('orgId');
+    expect(response).toEqual(testInvoiceSummary);
+  });
+
+  it('getInvoicePdf', () => {
+    const promise = subject().getInvoicePdf('id', 'orgId');
+
+    const response1 = new GetInvoicePdfResponse();
+    const chunk1 = new Uint8Array([1, 2]);
+    response1.setChunk(chunk1);
+    testGetInvoicePdfStream?.emit('data', response1);
+
+    const response2 = new GetInvoicePdfResponse();
+    const chunk2 = new Uint8Array([3, 4]);
+    response2.setChunk(chunk2);
+    testGetInvoicePdfStream?.emit('data', response2);
+    testGetInvoicePdfStream?.emit('end', { code: 0 });
+
+    const array = new Uint8Array([1, 2, 3, 4]);
+    expect(promise).resolves.toStrictEqual(array);
+  });
+});

src/app/billing-client.ts

@@ -0,0 +1,117 @@
+import { type RpcOptions } from '@improbable-eng/grpc-web/dist/typings/client.d';
+import { BillingServiceClient } from '../gen/app/v1/billing_pb_service';
+import pb from '../gen/app/v1/billing_pb';
+import { promisify } from '../utils';
+
+type GetCurrentMonthUsageResponse =
+  Partial<pb.GetCurrentMonthUsageResponse.AsObject> & {
+    start?: Date;
+    end?: Date;
+  };
+
+export class BillingClient {
+  private service: BillingServiceClient;
+
+  constructor(serviceHost: string, grpcOptions: RpcOptions) {
+    this.service = new BillingServiceClient(serviceHost, grpcOptions);
+  }
+
+  async getCurrentMonthUsage(orgId: string) {
+    const { service } = this;
+
+    const req = new pb.GetCurrentMonthUsageRequest();
+    req.setOrgId(orgId);
+
+    const response = await promisify<
+      pb.GetCurrentMonthUsageRequest,
+      pb.GetCurrentMonthUsageResponse
+    >(service.getCurrentMonthUsage.bind(service), req);
+
+    const result: GetCurrentMonthUsageResponse = response.toObject();
+    result.start = response.getStartDate()?.toDate();
+    result.end = response.getEndDate()?.toDate();
+    return result;
+  }
+
+  async getOrgBillingInformation(orgId: string) {
+    const { service } = this;
+
+    const req = new pb.GetOrgBillingInformationRequest();
+    req.setOrgId(orgId);
+
+    const response = await promisify<
+      pb.GetOrgBillingInformationRequest,
+      pb.GetOrgBillingInformationResponse
+    >(service.getOrgBillingInformation.bind(service), req);
+    return response.toObject();
+  }
+
+  async getInvoicesSummary(orgId: string) {
+    const { service } = this;
+
+    const req = new pb.GetInvoicesSummaryRequest();
+    req.setOrgId(orgId);
+
+    const response = await promisify<
+      pb.GetInvoicesSummaryRequest,
+      pb.GetInvoicesSummaryResponse
+    >(service.getInvoicesSummary.bind(service), req);
+    return response.toObject();
+  }
+
+  async getInvoicePdf(id: string, orgId: string) {
+    const { service } = this;
+
+    const req = new pb.GetInvoicePdfRequest();
+    req.setId(id);
+    req.setOrgId(orgId);
+
+    const chunks: Uint8Array[] = [];
+    const stream = service.getInvoicePdf(req);
+
+    stream.on('data', (response) => {
+      const chunk = response.getChunk_asU8();
+      chunks.push(chunk);
+    });
+
+    return new Promise<Uint8Array>((resolve, reject) => {
+      stream.on('status', (status) => {
+        if (status.code !== 0) {
+          const error = {
+            message: status.details,
+            code: status.code,
+            metadata: status.metadata,
+          };
+          reject(error);
+        }
+      });
+
+      stream.on('end', (end) => {
+        if (end === undefined) {
+          const error = { message: 'Stream ended without a status code' };
+          reject(error);
+        } else if (end.code !== 0) {
+          const error = {
+            message: end.details,
+            code: end.code,
+            metadata: end.metadata,
+          };
+          reject(error);
+        }
+        const arr = concatArrayU8(chunks);
+        resolve(arr);
+      });
+    });
+  }
+}
+
+const concatArrayU8 = (arrays: Uint8Array[]) => {
+  const totalLength = arrays.reduce((acc, value) => acc + value.length, 0);
+  const result = new Uint8Array(totalLength);
+  let length = 0;
+  for (const array of arrays) {
+    result.set(array, length);
+    length += array.length;
+  }
+  return result;
+};

src/app/viam-client.test.ts

@@ -13,6 +13,7 @@ vi.mock('./viam-transport', () => {
   };
 });
 import { DataClient } from './data-client';
+import { BillingClient } from './billing-client';
 import { createViamClient, type ViamClientOptions } from './viam-client';
 import { MlTrainingClient } from './ml-training-client';
 import { ProvisioningClient } from './provisioning-client';
@@ -47,6 +48,7 @@ describe('ViamClient', () => {
     expect(client.dataClient).toBeInstanceOf(DataClient);
     expect(client.mlTrainingClient).toBeInstanceOf(MlTrainingClient);
     expect(client.provisioningClient).toBeInstanceOf(ProvisioningClient);
+    expect(client.billingClient).toBeInstanceOf(BillingClient);
   });
 
   it('create client with an api key credential and a custom service host', async () => {
@@ -61,6 +63,7 @@ describe('ViamClient', () => {
     expect(client.dataClient).toBeInstanceOf(DataClient);
     expect(client.mlTrainingClient).toBeInstanceOf(MlTrainingClient);
     expect(client.provisioningClient).toBeInstanceOf(ProvisioningClient);
+    expect(client.billingClient).toBeInstanceOf(BillingClient);
   });
 
   it('create client with an access token', async () => {
@@ -74,5 +77,6 @@ describe('ViamClient', () => {
     expect(client.dataClient).toBeInstanceOf(DataClient);
     expect(client.mlTrainingClient).toBeInstanceOf(MlTrainingClient);
     expect(client.provisioningClient).toBeInstanceOf(ProvisioningClient);
+    expect(client.billingClient).toBeInstanceOf(BillingClient);
   });
 });

src/app/viam-client.ts

@@ -5,6 +5,7 @@ import {
   type AccessToken,
 } from './viam-transport';
 import { DataClient } from './data-client';
+import { BillingClient } from './billing-client';
 import { MlTrainingClient } from './ml-training-client';
 import { ProvisioningClient } from './provisioning-client';
 
@@ -35,6 +36,7 @@ export class ViamClient {
   public dataClient: DataClient | undefined;
   public mlTrainingClient: MlTrainingClient | undefined;
   public provisioningClient: ProvisioningClient | undefined;
+  public billingClient: BillingClient | undefined;
 
   constructor(transportFactory: grpc.TransportFactory, serviceHost: string) {
     this.transportFactory = transportFactory;
@@ -49,5 +51,6 @@ export class ViamClient {
       this.serviceHost,
       grpcOptions
     );
+    this.billingClient = new BillingClient(this.serviceHost, grpcOptions);
   }
 }

src/main.ts

@@ -86,6 +86,18 @@ export {
   type ProvisioningClient,
 } from './app/provisioning-client';
 
+/**
+ * Raw Protobuf interfaces for Billing.
+ *
+ * Generated with https://github.com/improbable-eng/grpc-web
+ *
+ * @deprecated Use {@link BillingClient} instead.
+ * @alpha
+ * @group Raw Protobufs
+ */
+export { default as billingApi } from './gen/app/v1/billing_pb';
+export { type BillingClient } from './app/billing-client';
+
 /**
  * Raw Protobuf interfaces for an Arm component.
  *