From fb49474ac96163f819f3bc647b599eb7ccdd60b6 Mon Sep 17 00:00:00 2001
From: Daniel Rocha <daniel.rocha@consensys.net>
Date: Fri, 7 Jun 2024 17:44:50 +0200
Subject: [PATCH] feat: add `getAccountBalances` to `KeyringClient` (#340)

---
 src/KeyringClient.test.ts | 58 +++++++++++++++++++++++++++++++++++++++
 src/KeyringClient.ts      | 16 +++++++++++
 2 files changed, 74 insertions(+)

diff --git a/src/KeyringClient.test.ts b/src/KeyringClient.test.ts
index 62ba2621b..c5d237b74 100644
--- a/src/KeyringClient.test.ts
+++ b/src/KeyringClient.test.ts
@@ -3,6 +3,7 @@ import {
   type KeyringRequest,
   type KeyringResponse,
   KeyringClient,
+  KeyringRpcMethod,
 } from '.'; // Import from `index.ts` to test the public API
 
 describe('KeyringClient', () => {
@@ -84,6 +85,63 @@ describe('KeyringClient', () => {
     });
   });
 
+  describe('getAccountBalances', () => {
+    it('returns a valid response', async () => {
+      const assets = ['bip122:000000000019d6689c085ae165831e93/slip44:0'];
+      const id = '1617ea08-d4b6-48bf-ba83-901ef1e45ed7';
+      const expectedResponse = {
+        [assets[0] as string]: {
+          amount: '1234',
+          unit: 'sat',
+        },
+      };
+
+      mockSender.send.mockResolvedValue(expectedResponse);
+      const balances = await keyring.getAccountBalances(id, assets);
+
+      expect(mockSender.send).toHaveBeenCalledWith({
+        jsonrpc: '2.0',
+        id: expect.any(String),
+        method: `${KeyringRpcMethod.GetAccountBalances}`,
+        params: { id, assets },
+      });
+
+      expect(balances).toStrictEqual(expectedResponse);
+    });
+
+    it('throws an error because the amount has the wrong type', async () => {
+      const assets = ['bip122:000000000019d6689c085ae165831e93/slip44:0'];
+      const id = '1617ea08-d4b6-48bf-ba83-901ef1e45ed7';
+      const expectedResponse = {
+        [assets[0] as string]: {
+          amount: 1234, // Should be a `StringNumber`
+          unit: 'sat',
+        },
+      };
+
+      mockSender.send.mockResolvedValue(expectedResponse);
+      await expect(keyring.getAccountBalances(id, assets)).rejects.toThrow(
+        'At path: bip122:000000000019d6689c085ae165831e93/slip44:0.amount -- Expected a value of type `StringNumber`, but received: `1234`',
+      );
+    });
+
+    it("throws an error because the amount isn't a StringNumber", async () => {
+      const assets = ['bip122:000000000019d6689c085ae165831e93/slip44:0'];
+      const id = '1617ea08-d4b6-48bf-ba83-901ef1e45ed7';
+      const expectedResponse = {
+        [assets[0] as string]: {
+          amount: 'not-a-string-number', // Should be a `StringNumber`
+          unit: 'sat',
+        },
+      };
+
+      mockSender.send.mockResolvedValue(expectedResponse);
+      await expect(keyring.getAccountBalances(id, assets)).rejects.toThrow(
+        'At path: bip122:000000000019d6689c085ae165831e93/slip44:0.amount -- Expected a value of type `StringNumber`, but received: `"not-a-string-number"`',
+      );
+    });
+  });
+
   describe('filterAccountChains', () => {
     it('should send a request to filter the chains supported by an account and return the response', async () => {
       const id = '49116980-0712-4fa5-b045-e4294f1d440e';
diff --git a/src/KeyringClient.ts b/src/KeyringClient.ts
index 0f4f1923f..9b6b16c5d 100644
--- a/src/KeyringClient.ts
+++ b/src/KeyringClient.ts
@@ -8,6 +8,8 @@ import type {
   KeyringRequest,
   KeyringAccountData,
   KeyringResponse,
+  CaipAssetType,
+  Balance,
 } from './api';
 import {
   ApproveRequestResponseStruct,
@@ -15,6 +17,7 @@ import {
   DeleteAccountResponseStruct,
   ExportAccountResponseStruct,
   FilterAccountChainsResponseStruct,
+  GetAccountBalancesResponseStruct,
   GetAccountResponseStruct,
   GetRequestResponseStruct,
   ListAccountsResponseStruct,
@@ -76,6 +79,19 @@ export class KeyringClient implements Keyring {
     );
   }
 
+  async getAccountBalances(
+    id: string,
+    assets: CaipAssetType[],
+  ): Promise<Record<CaipAssetType, Balance>> {
+    return strictMask(
+      await this.#send({
+        method: KeyringRpcMethod.GetAccountBalances,
+        params: { id, assets },
+      }),
+      GetAccountBalancesResponseStruct,
+    );
+  }
+
   async createAccount(
     options: Record<string, Json> = {},
   ): Promise<KeyringAccount> {