fix(sdk): Fix Bifrost balance fetching 🔧

packages/sdk/src/api/IPolkadotApi.ts

@@ -21,6 +21,7 @@ export interface IPolkadotApi<TApi, TRes> {
   getAssetHubForeignBalance(address: string, multiLocation: TMultiLocation): Promise<bigint>
   getForeignAssetsByIdBalance(address: string, assetId: string): Promise<bigint>
   getBalanceForeignXTokens(address: string, asset: TAsset): Promise<bigint>
+  getBalanceForeignBifrost(address: string, asset: TAsset): Promise<bigint>
   getBalanceForeignAssetsAccount(address: string, assetId: bigint | number): Promise<bigint>
   getFromStorage(key: string): Promise<string>
   clone(): IPolkadotApi<TApi, TRes>

packages/sdk/src/nodes/supported/BifrostPolkadot.ts

@@ -25,7 +25,7 @@ export class BifrostPolkadot<TApi, TRes>
     super('BifrostPolkadot', 'bifrost', 'polkadot', Version.V3)
   }
 
-  private getCurrencySelection(asset: TAsset) {
+  getCurrencySelection(asset: TAsset) {
     const nativeAssetSymbol = this.getNativeAssetSymbol()
 
     if (asset.symbol === nativeAssetSymbol) {

packages/sdk/src/pallets/assets/balance/getBalanceForeignXTokens.test.ts

@@ -7,7 +7,8 @@ import { getBalanceForeignXTokens } from './getBalanceForeignXTokens'
 
 const mockApi = {
   getBalanceForeignAssetsAccount: vi.fn(),
-  getBalanceForeignXTokens: vi.fn()
+  getBalanceForeignXTokens: vi.fn(),
+  getBalanceForeignBifrost: vi.fn()
 } as unknown as IPolkadotApi<ApiPromise, Extrinsic>
 
 describe('getBalanceForeignXTokens', () => {
@@ -45,6 +46,15 @@ describe('getBalanceForeignXTokens', () => {
     expect(balance).toBe(BigInt(3000))
   })
 
+  it('calls getBalanceForeignBifrost when node is BifrostPolkadot', async () => {
+    const spy = vi.spyOn(mockApi, 'getBalanceForeignBifrost').mockResolvedValue(BigInt(4000))
+    const spy2 = vi.spyOn(mockApi, 'getBalanceForeignXTokens')
+    const balance = await getBalanceForeignXTokens(mockApi, 'BifrostPolkadot', address, asset)
+    expect(spy).toHaveBeenCalledWith(address, asset)
+    expect(spy2).not.toHaveBeenCalled()
+    expect(balance).toBe(BigInt(4000))
+  })
+
   it('returns 0 if getBalanceForeignMoonbeam resolves with 0', async () => {
     vi.spyOn(mockApi, 'getBalanceForeignAssetsAccount').mockResolvedValue(BigInt(0))
     const balance = await getBalanceForeignXTokens(mockApi, 'Moonbeam', address, asset)

packages/sdk/src/pallets/assets/balance/getBalanceForeignXTokens.ts

@@ -17,5 +17,9 @@ export const getBalanceForeignXTokens = async <TApi, TRes>(
     return api.getBalanceForeignAssetsAccount(address, BigInt(asset.assetId))
   }
 
+  if (node === 'BifrostPolkadot' || node === 'BifrostKusama') {
+    return api.getBalanceForeignBifrost(address, asset)
+  }
+
   return api.getBalanceForeignXTokens(address, asset)
 }

packages/sdk/src/papi/PapiApi.test.ts

@@ -28,7 +28,7 @@ vi.mock('./PapiXcmTransformer', () => ({
   transform: vi.fn().mockReturnValue({ transformed: true })
 }))
 
-vi.mock('../utils', () => ({
+vi.mock('../utils/createApiInstanceForNode', () => ({
   createApiInstanceForNode: vi.fn().mockResolvedValue({} as PolkadotClient)
 }))
 
@@ -85,6 +85,9 @@ describe('PapiApi', () => {
           },
           Tokens: {
             Accounts: {
+              getValue: vi.fn().mockResolvedValue({
+                free: BigInt(6000)
+              }),
               getEntries: vi.fn().mockResolvedValue([
                 {
                   keyArgs: [
@@ -289,9 +292,30 @@ describe('PapiApi', () => {
     })
   })
 
+  describe('getBalanceForeignBifrost', () => {
+    it('should return the balance when balance exists', async () => {
+      const transformedCurrencySelection = { type: 'Native', value: 'BNC' }
+
+      vi.mocked(transform).mockReturnValue(transformedCurrencySelection)
+
+      papiApi.setApi(mockPolkadotClient)
+      const balance = await papiApi.getBalanceForeignBifrost('some_address', {
+        symbol: 'BNC'
+      })
+
+      const unsafeApi = papiApi.getApi().getUnsafeApi()
+      expect(unsafeApi.query.Tokens.Accounts.getValue).toHaveBeenCalledWith(
+        'some_address',
+        transformedCurrencySelection
+      )
+      expect(balance).toBe(BigInt(6000))
+    })
+  })
+
   describe('getBalanceForeignXTokens', () => {
     it('should return the balance when asset matches symbolOrId', async () => {
       papiApi.setApi(mockPolkadotClient)
+
       const balance = await papiApi.getBalanceForeignXTokens('some_address', {
         symbol: 'DOT',
         assetId: '1'

packages/sdk/src/papi/PapiApi.ts

@@ -1,3 +1,5 @@
+/* eslint-disable @typescript-eslint/no-unsafe-argument */
+/* eslint-disable @typescript-eslint/no-unsafe-assignment */
 /* eslint-disable @typescript-eslint/no-unsafe-member-access */
 /* eslint-disable @typescript-eslint/no-unsafe-call */
 import type {
@@ -8,7 +10,7 @@ import type {
   TNodeWithRelayChains,
   TSerializedApiCallV2
 } from '../types'
-import { createApiInstanceForNode } from '../utils'
+import { createApiInstanceForNode, getNode } from '../utils'
 import { createClient, FixedSizeBinary } from 'polkadot-api'
 import type { TPapiApi, TPapiApiOrUrl, TPapiTransaction } from '../papi/types'
 import type { IPolkadotApi } from '../api'
@@ -73,10 +75,8 @@ class PapiApi implements IPolkadotApi<TPapiApi, TPapiTransaction> {
   }
 
   callTxMethod({ module, section, parameters }: TSerializedApiCallV2): TPapiTransaction {
-    // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
     const transformedParameters = transform(parameters)
 
-    // eslint-disable-next-line @typescript-eslint/no-unsafe-argument
     return this.api.getUnsafeApi().tx[module][section](transformedParameters)
   }
 
@@ -85,53 +85,56 @@ class PapiApi implements IPolkadotApi<TPapiApi, TPapiTransaction> {
   }
 
   async getBalanceNative(address: string): Promise<bigint> {
-    // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
     const res = await this.api.getUnsafeApi().query.System.Account.getValue(address)
 
     return res.data.free as bigint
   }
 
   async getBalanceForeignPolkadotXcm(address: string, id?: string): Promise<bigint> {
-    // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
     const res = await this.api.getUnsafeApi().query.Assets.Account.getValue(id, address)
 
-    // eslint-disable-next-line @typescript-eslint/no-unsafe-argument
     return res && res.balance ? BigInt(res.balance) : BigInt(0)
   }
 
   async getMythosForeignBalance(address: string): Promise<bigint> {
-    // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
     const res = await this.api.getUnsafeApi().query.Balances.Account.getValue(address)
-    // eslint-disable-next-line @typescript-eslint/no-unsafe-argument
+
     return res && res.free ? BigInt(res.free) : BigInt(0)
   }
 
   async getAssetHubForeignBalance(address: string, multiLocation: TMultiLocation): Promise<bigint> {
-    // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
     const transformedMultiLocation = transform(multiLocation)
 
-    // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
     const res = await this.api
       .getUnsafeApi()
       .query.ForeignAssets.Account.getValue(transformedMultiLocation, address)
 
-    // eslint-disable-next-line @typescript-eslint/no-unsafe-argument
     return BigInt(res === undefined ? 0 : res.balance)
   }
 
   async getForeignAssetsByIdBalance(address: string, assetId: string): Promise<bigint> {
-    // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
     const res = await this.api.getUnsafeApi().query.ForeignAssets.Account.getValue(assetId, address)
 
-    // eslint-disable-next-line @typescript-eslint/no-unsafe-argument
     return BigInt(res === undefined ? 0 : res.balance)
   }
 
+  async getBalanceForeignBifrost(address: string, asset: TAsset): Promise<bigint> {
+    const currencySelection = getNode('BifrostPolkadot').getCurrencySelection(asset)
+
+    const transformedParameters = transform(currencySelection)
+
+    const response = await this.api
+      .getUnsafeApi()
+      .query.Tokens.Accounts.getValue(address, transformedParameters)
+
+    const accountData = response ? response : null
+    return accountData ? BigInt(accountData.free.toString()) : BigInt(0)
+  }
+
   async getBalanceForeignXTokens(address: string, asset: TAsset): Promise<bigint> {
     const response = await this.api.getUnsafeApi().query.Tokens.Accounts.getEntries(address)
 
     const entry = response.find(({ keyArgs }) => {
-      // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
       const [_address, assetItem] = keyArgs
 
       return (
@@ -147,14 +150,13 @@ class PapiApi implements IPolkadotApi<TPapiApi, TPapiTransaction> {
           assetItem.value.toString().toLowerCase() === asset.assetId?.toLowerCase())
       )
     })
-    // eslint-disable-next-line @typescript-eslint/no-unsafe-argument
+
     return entry?.value ? BigInt(entry.value.free.toString()) : BigInt(0)
   }
 
   async getBalanceForeignAssetsAccount(address: string, assetId: bigint | number): Promise<bigint> {
-    // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
     const response = await this.api.getUnsafeApi().query.Assets.Account.getValue(assetId, address)
-    // eslint-disable-next-line @typescript-eslint/no-unsafe-argument
+
     return BigInt(response === undefined ? 0 : response.balance)
   }
 

packages/sdk/src/pjs/PolkadotJsApi.test.ts

@@ -48,9 +48,9 @@ describe('PolkadotJsApi', () => {
           account: vi.fn()
         },
         tokens: {
-          accounts: {
+          accounts: Object.assign(vi.fn(), {
             entries: vi.fn()
-          }
+          })
         }
       }
     } as unknown as TPjsApi
@@ -277,6 +277,41 @@ describe('PolkadotJsApi', () => {
     })
   })
 
+  describe('getBalanceForeignBifrost', () => {
+    it('should return the balance when asset matches currencySelection', async () => {
+      const address = 'some_address'
+      const mockResponse = {
+        free: { toString: () => '6000' }
+      } as unknown as VoidFn
+
+      vi.mocked(mockApiPromise.query.tokens.accounts).mockResolvedValue(mockResponse)
+
+      const balance = await polkadotApi.getBalanceForeignBifrost(address, { symbol: 'DOT' })
+
+      expect(mockApiPromise.query.tokens.accounts).toHaveBeenCalledWith(address, {
+        Token: 'DOT'
+      })
+      expect(balance).toBe(BigInt(6000))
+    })
+
+    it('should return null when no matching asset found', async () => {
+      const address = 'some_address'
+
+      const mockResponse = {
+        free: { toString: () => '0' }
+      } as unknown as VoidFn
+
+      vi.mocked(mockApiPromise.query.tokens.accounts).mockResolvedValue(mockResponse)
+
+      const balance = await polkadotApi.getBalanceForeignBifrost(address, { symbol: 'DOT' })
+
+      expect(mockApiPromise.query.tokens.accounts).toHaveBeenCalledWith(address, {
+        Token: 'DOT'
+      })
+      expect(balance).toBe(BigInt(0))
+    })
+  })
+
   describe('getBalanceForeignXTokens', () => {
     it('should return the balance when asset matches symbolOrId', async () => {
       const address = 'some_address'

packages/sdk/src/pjs/PolkadotJsApi.ts

@@ -13,7 +13,7 @@ import type { StorageKey } from '@polkadot/types'
 import { u32, type UInt } from '@polkadot/types'
 import type { AnyTuple, Codec } from '@polkadot/types/types'
 import type { TBalanceResponse } from '../types/TBalance'
-import { createApiInstanceForNode } from '../utils'
+import { createApiInstanceForNode, getNode } from '../utils'
 import { isForeignAsset } from '../utils/assets'
 
 const lowercaseFirstLetter = (value: string) => value.charAt(0).toLowerCase() + value.slice(1)
@@ -99,6 +99,15 @@ class PolkadotJsApi implements IPolkadotApi<TPjsApi, Extrinsic> {
     return BigInt(obj === null || !obj.balance ? 0 : obj.balance)
   }
 
+  async getBalanceForeignBifrost(address: string, asset: TAsset): Promise<bigint> {
+    const currencySelection = getNode('BifrostPolkadot').getCurrencySelection(asset)
+
+    const response: Codec = await this.api.query.tokens.accounts(address, currencySelection)
+
+    const accountData = response ? (response as AccountData) : null
+    return accountData ? BigInt(accountData.free.toString()) : BigInt(0)
+  }
+
   async getBalanceForeignXTokens(address: string, asset: TAsset): Promise<bigint> {
     const response: Array<[StorageKey<AnyTuple>, Codec]> =
       await this.api.query.tokens.accounts.entries(address)