Feature: extendable setup to define initial token balances when rivet starts

[agualis]

I use Rivet to manual test new features in my Dapp.
I normally use the default rivet accounts to make things simpler but I need to manually import some tokens and manually add balances for them (the tokens and balances required by my feature under test).

I crafted this ugly useInitialTokenBalances.ts hook to perform that task:

import { useEffect } from 'react'
import { parseEther, type Address } from 'viem'
import { useAccountTokens } from './useAccountTokens'
import { useSetErc20Balance } from './useSetErc20Balance'

type TokenBalancesInput = Record<Address, bigint>

const input: TokenBalancesInput = {
  '0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2': parseEther('100'),
  '0x198d7387fa97a73f05b8578cdeff8f2a1f34cd1f': parseEther('100'),
}

export function useInitialTokenBalances(initialTokenBalances: TokenBalancesInput = input) {
  // const { account } = useAccount()
  const accountAddress = '0xf39Fd6e51aad88F6F4ce6aB8827279cffFb92266'
  const { addToken, tokens, isLoading } = useAccountTokens({ address: accountAddress })
  const { mutate } = useSetErc20Balance()

  useEffect(()=> {
    if (!isLoading) {
      Object.keys(initialTokenBalances).forEach((tokenAddress) => {
        const balance = initialTokenBalances[tokenAddress as Address]

        // Only set balance if the token is not already present
        if (!tokens?.map(t=> t.address).includes(tokenAddress)) {

          console.log(`Adding default token balance for: ${tokenAddress}`)

          addToken({tokenAddress: tokenAddress as Address})

          mutate({
            address: accountAddress,
            tokenAddress: tokenAddress as Address,
            value: balance,
          })
        }
      })
    }
  }, [isLoading])
}

so that I can call it from, for instance, src/screens/index.tsx.

I think that it could be worthy to clean the hook and adding an extendable system to allow users to define initial token state from a data structure similar to this one:

const initialBalances: TokenBalancesInput = {
  '0xc02aaa39b223fe8d0a0e5c4f27ead9083c756cc2': parseEther('100'),
  '0x198d7387fa97a73f05b8578cdeff8f2a1f34cd1f': parseEther('100'),
}