From 9aaa96aec9679d7f90b36b1484835871f1956d0e Mon Sep 17 00:00:00 2001
From: "juraj.bacovcin" <juraj.bacovcin@tatum.io>
Date: Wed, 10 Apr 2024 12:31:25 +0200
Subject: [PATCH] ALL-6149 Add tokenTypes filter to address getBalance

---
 CHANGELOG.md                   |  6 ++++++
 package.json                   |  2 +-
 src/dto/shared.dto.ts          |  9 ++++++++-
 src/e2e/tatum.address.spec.ts  | 15 +++++++++++++++
 src/service/address/address.ts | 29 ++++++++++++++++++++++-------
 5 files changed, 52 insertions(+), 9 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7aa824023..6fcad452c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,9 @@
+## [4.2.23] - 2024.4.10
+
+### Updated
+
+- Improved `getBalance` method in `address` module by introducing `tokenTypes` filter along with `native` option.
+
 ## [4.2.22] - 2024.4.10
 
 ### Fixed
diff --git a/package.json b/package.json
index cc8679296..e1f7accf9 100644
--- a/package.json
+++ b/package.json
@@ -1,6 +1,6 @@
 {
   "name": "@tatumio/tatum",
-  "version": "4.2.22",
+  "version": "4.2.23",
   "description": "Tatum JS SDK",
   "author": "Tatum",
   "repository": "https://github.com/tatumio/tatum-js",
diff --git a/src/dto/shared.dto.ts b/src/dto/shared.dto.ts
index eb4514353..0f70f71d0 100644
--- a/src/dto/shared.dto.ts
+++ b/src/dto/shared.dto.ts
@@ -32,6 +32,12 @@ export interface AddressBalanceFilters extends Pagination {
    * List of addresses to check.
    */
   addresses: string[]
+
+  /**
+   * Optional filter for token types. If not specified, all token types are returned.
+   * Allowed values are `native`, `fungible`, `nft` and `multitoken`.
+   */
+  tokenTypes?: string[]
 }
 
 export interface AddressBalanceFiltersTron {
@@ -48,7 +54,8 @@ export interface AddressBalanceFiltersTezos extends Pagination {
   address: string
 
   /**
-   * Optional filter for token types. If not specified, all token types are returned. Allowed values are `fungible`, `nft` and `multitoken`.
+   * Optional filter for token types. If not specified, all token types are returned.
+   * Allowed values are `fungible`, `nft` and `multitoken`.
    */
   tokenTypes?: string[]
 }
diff --git a/src/e2e/tatum.address.spec.ts b/src/e2e/tatum.address.spec.ts
index 3f52f579b..47cbfb41b 100644
--- a/src/e2e/tatum.address.spec.ts
+++ b/src/e2e/tatum.address.spec.ts
@@ -45,6 +45,21 @@ describe.skip('Address', () => {
         })
       })
 
+      it('should get only native balance with native assets only', async () => {
+        const { data } = await tatum.address.getBalance({
+          addresses: ['0x514D547c8aC8ccBEc29b5144810454BD7d3625CA'],
+          tokenTypes: ['native'],
+        })
+        expect(data).toHaveLength(1)
+        expect(data[0]).toStrictEqual({
+          asset: 'ETH',
+          decimals: 18,
+          address: '0x514D547c8aC8ccBEc29b5144810454BD7d3625CA',
+          balance: expect.any(String),
+          type: 'native',
+        })
+      })
+
       it('should get balance with native assets only for 2 addresses', async () => {
         const { data } = await tatum.address.getBalance({
           addresses: [
diff --git a/src/service/address/address.ts b/src/service/address/address.ts
index 5d7d9ead2..3a6496d7b 100644
--- a/src/service/address/address.ts
+++ b/src/service/address/address.ts
@@ -246,10 +246,25 @@ export class Address {
     page = 0,
     pageSize = 10,
     addresses,
+    tokenTypes,
   }: AddressBalanceFilters): Promise<ResponseDto<AddressBalance[]>> {
     const chain = this.config.network
     return ErrorUtils.tryFail(async () => {
-      const nativeBalances = await this.getNativeBalance(addresses)
+      const result: AddressBalance[] = []
+
+      if (!tokenTypes || tokenTypes.includes('native')) {
+        const nativeBalances = await this.getNativeBalance(addresses)
+        result.push(...formatNativeBalances(nativeBalances, addresses, chain))
+
+        if (tokenTypes) {
+          tokenTypes = tokenTypes.filter((tokenType) => tokenType !== 'native')
+        }
+      }
+
+      if (tokenTypes?.length === 0) {
+        return result
+      }
+
       const tokenBalances =
         isDataApiEvmEnabledNetwork(chain) &&
         (await this.connector
@@ -261,17 +276,17 @@ export class Address {
               excludeMetadata: true,
               chain,
               addresses: addresses.join(','),
+              tokenTypes: tokenTypes?.join(','),
             },
           })
           .then((r) => r.result))
 
-      const result = formatNativeBalances(nativeBalances, addresses, chain)
-
-      if (!tokenBalances) {
-        return result
+      if (tokenBalances) {
+        const serializedTokenBalances = await this.processTokenBalanceDetails(tokenBalances, chain)
+        result.push(...serializedTokenBalances)
       }
-      const serializedTokenBalances = await this.processTokenBalanceDetails(tokenBalances, chain)
-      return [...result, ...serializedTokenBalances]
+
+      return result
     })
   }