Introducing AssetBalance type

api/lib/src/queries.rs

@@ -116,6 +116,11 @@ impl QueryApi {
 		}))
 		.await
 		.into_iter()
+		.map(|result| {
+			result.map(|(asset, maybe_balance)| {
+				(asset, if let Some(balance) = maybe_balance { balance.amount() } else { 0u128 })
+			})
+		})
 		.collect()
 	}
 

bouncer/shared/get_free_balance.ts

@@ -0,0 +1,12 @@
+import { InternalAsset as Asset } from '@chainflip/cli';
+import { getChainflipApi } from './utils/substrate';
+
+export async function getFreeBalance(address: string, asset: Asset): Promise<bigint> {
+  await using chainflip = await getChainflipApi();
+  const fee = await chainflip.query.assetBalances.freeBalances(address, asset);
+  // If the option is none we assume the balance is 0 for tests.
+  if (fee.isEmpty) {
+    return BigInt(0);
+  }
+  return BigInt(JSON.parse(fee.toString()).amount);
+}

bouncer/tests/broker_fee_collection.ts

@@ -1,8 +1,6 @@
 import assert from 'assert';
 import { randomBytes } from 'crypto';
 import { InternalAsset as Asset, InternalAssets as Assets } from '@chainflip/cli';
-// eslint-disable-next-line no-restricted-imports
-import { KeyringPair } from '@polkadot/keyring/types';
 import Keyring from '../polkadot/keyring';
 import {
   brokerMutex,
@@ -20,6 +18,7 @@ import { getBalance } from '../shared/get_balance';
 import { getChainflipApi, observeEvent } from '../shared/utils/substrate';
 import { send } from '../shared/send';
 import { ExecutableTest } from '../shared/executable_test';
+import { getFreeBalance } from '../shared/get_free_balance';
 
 /* eslint-disable @typescript-eslint/no-use-before-define */
 export const testBrokerFeeCollection = new ExecutableTest('Broker-Fee-Collection', main, 200);
@@ -56,21 +55,12 @@ export async function submitBrokerWithdrawal(
 
 const feeAsset = Assets.Usdc;
 
-async function getEarnedBrokerFees(brokerKeypair: KeyringPair): Promise<bigint> {
-  await using chainflip = await getChainflipApi();
-  // NOTE: All broker fees are collected in USDC now:
-  const feeStr = (
-    await chainflip.query.assetBalances.freeBalances(brokerKeypair.address, Assets.Usdc)
-  ).toString();
-  return BigInt(feeStr);
-}
-
 /// Runs a swap, checks that the broker fees are collected,
 /// then withdraws the broker fees, making sure the balance is correct after the withdrawal.
 async function testBrokerFees(inputAsset: Asset, seed?: string): Promise<void> {
   await using chainflip = await getChainflipApi();
   // Check the broker fees before the swap
-  const earnedBrokerFeesBefore = await getEarnedBrokerFees(broker);
+  const earnedBrokerFeesBefore = await getFreeBalance(broker.address, Assets.Usdc);
   testBrokerFeeCollection.log(`${inputAsset} earnedBrokerFeesBefore:`, earnedBrokerFeesBefore);
 
   // Run a swap
@@ -153,7 +143,7 @@ async function testBrokerFees(inputAsset: Asset, seed?: string): Promise<void> {
   );
 
   // Check that the detected increase in earned broker fees matches the swap event values and it is equal to the expected amount (after the deposit fee is accounted for)
-  const earnedBrokerFeesAfter = await getEarnedBrokerFees(broker);
+  const earnedBrokerFeesAfter = await getFreeBalance(broker.address, Assets.Usdc);
   testBrokerFeeCollection.log(`${inputAsset} earnedBrokerFeesAfter:`, earnedBrokerFeesAfter);
 
   assert(earnedBrokerFeesAfter > earnedBrokerFeesBefore, 'No increase in earned broker fees');

bouncer/tests/lp_api_test.ts

@@ -17,8 +17,9 @@ import { jsonRpc } from '../shared/json_rpc';
 import { depositLiquidity } from '../shared/deposit_liquidity';
 import { sendEvmNative } from '../shared/send_evm';
 import { getBalance } from '../shared/get_balance';
-import { getChainflipApi, observeEvent } from '../shared/utils/substrate';
+import { observeEvent } from '../shared/utils/substrate';
 import { ExecutableTest } from '../shared/executable_test';
+import { getFreeBalance } from '../shared/get_free_balance';
 
 /* eslint-disable @typescript-eslint/no-use-before-define */
 export const testLpApi = new ExecutableTest('LP-API', main, 200);
@@ -145,29 +146,24 @@ async function testWithdrawAsset() {
 }
 
 async function testTransferAsset() {
-  await using chainflip = await getChainflipApi();
   const amountToTransfer = testAssetAmount.toString(16);
 
-  const getLpBalance = async (account: string) =>
-    // eslint-disable-next-line @typescript-eslint/no-explicit-any
-    ((await chainflip.query.assetBalances.freeBalances(account, testAsset)) as any).toBigInt();
-
   const keyring = new Keyring({ type: 'sr25519' });
 
   const sourceLpAccount = keyring.createFromUri('//LP_1');
   const destinationLpAccount = keyring.createFromUri('//LP_2');
 
-  const oldBalanceSource = await getLpBalance(sourceLpAccount.address);
-  const oldBalanceDestination = await getLpBalance(destinationLpAccount.address);
+  const oldBalanceSource = await getFreeBalance(sourceLpAccount.address, testAsset);
+  const oldBalanceDestination = await getFreeBalance(destinationLpAccount.address, testAsset);
 
   const result = await lpApiRpc(`lp_transfer_asset`, [
     amountToTransfer,
     testRpcAsset,
     destinationLpAccount.address,
   ]);
 
-  let newBalancesSource = await getLpBalance(sourceLpAccount.address);
-  let newBalanceDestination = await getLpBalance(destinationLpAccount.address);
+  let newBalancesSource = await getFreeBalance(sourceLpAccount.address, testAsset);
+  let newBalanceDestination = await getFreeBalance(destinationLpAccount.address, testAsset);
 
   // Wait max for 18 seconds aka 3 blocks for the balances to update.
   for (let i = 0; i < 18; i++) {
@@ -177,8 +173,8 @@ async function testTransferAsset() {
 
     await sleep(1000);
 
-    newBalancesSource = await getLpBalance(sourceLpAccount.address);
-    newBalanceDestination = await getLpBalance(destinationLpAccount.address);
+    newBalancesSource = await getFreeBalance(sourceLpAccount.address, testAsset);
+    newBalanceDestination = await getFreeBalance(destinationLpAccount.address, testAsset);
   }
 
   // Expect result to be a block hash

state-chain/cf-integration-tests/src/swapping.rs

@@ -74,10 +74,18 @@ fn new_pool(unstable_asset: Asset, fee_hundredth_pips: u32, initial_price: Price
 }
 
 fn credit_account(account_id: &AccountId, asset: Asset, amount: AssetAmount) {
-	let original_amount = pallet_cf_asset_balances::FreeBalances::<Runtime>::get(account_id, asset);
+	let original_amount = if let Some(balance) =
+		pallet_cf_asset_balances::FreeBalances::<Runtime>::get(account_id, asset)
+	{
+		balance.amount()
+	} else {
+		0u128
+	};
 	assert_ok!(AssetBalances::try_credit_account(account_id, asset, amount));
 	assert_eq!(
-		pallet_cf_asset_balances::FreeBalances::<Runtime>::get(account_id, asset),
+		pallet_cf_asset_balances::FreeBalances::<Runtime>::get(account_id, asset)
+			.unwrap()
+			.amount(),
 		original_amount + amount
 	);
 	assert_has_matching_event!(
@@ -114,11 +122,15 @@ fn set_range_order(
 	let new_balances = [base_asset, quote_asset]
 		.map(|asset| pallet_cf_asset_balances::FreeBalances::<Runtime>::get(account_id, asset));
 
-	assert!(new_balances.into_iter().zip(balances).all(|(new, old)| { new <= old }));
+	assert!(new_balances.iter().zip(balances.as_ref()).all(|(new, old)| { new <= old }));
 
-	for ((new_balance, old_balance), expected_asset) in
-		new_balances.into_iter().zip(balances).zip([base_asset, quote_asset])
+	for ((maybe_new_balance, maybe_old_balance), expected_asset) in
+		new_balances.iter().zip(balances.as_ref()).zip([base_asset, quote_asset])
 	{
+		let (new_balance, old_balance) = (
+			maybe_new_balance.as_ref().unwrap().amount(),
+			maybe_old_balance.as_ref().unwrap().amount(),
+		);
 		if new_balance < old_balance {
 			assert_has_matching_event!(
 				Runtime,
@@ -146,8 +158,12 @@ fn set_limit_order(
 	let (asset_pair, order) = pallet_cf_pools::AssetPair::from_swap(sell_asset, buy_asset).unwrap();
 
 	let sell_balance =
-		pallet_cf_asset_balances::FreeBalances::<Runtime>::get(account_id, sell_asset);
-	let buy_balance = pallet_cf_asset_balances::FreeBalances::<Runtime>::get(account_id, buy_asset);
+		pallet_cf_asset_balances::FreeBalances::<Runtime>::get(account_id, sell_asset)
+			.unwrap()
+			.amount();
+	let buy_balance = pallet_cf_asset_balances::FreeBalances::<Runtime>::get(account_id, buy_asset)
+		.unwrap()
+		.amount();
 	assert_ok!(LiquidityPools::set_limit_order(
 		RuntimeOrigin::signed(account_id.clone()),
 		asset_pair.assets().base,
@@ -158,9 +174,13 @@ fn set_limit_order(
 		sell_amount,
 	));
 	let new_sell_balance =
-		pallet_cf_asset_balances::FreeBalances::<Runtime>::get(account_id, sell_asset);
+		pallet_cf_asset_balances::FreeBalances::<Runtime>::get(account_id, sell_asset)
+			.unwrap()
+			.amount();
 	let new_buy_balance =
-		pallet_cf_asset_balances::FreeBalances::<Runtime>::get(account_id, buy_asset);
+		pallet_cf_asset_balances::FreeBalances::<Runtime>::get(account_id, buy_asset)
+			.unwrap()
+			.amount();
 
 	assert_eq!(new_sell_balance, sell_balance - sell_amount);
 	assert_eq!(new_buy_balance, buy_balance);