Override ERC-20 balance for an address using Viem & simulateContract

overrideBalance.ts

import { createPublicClient, keccak256, numberToHex, pad, toHex, erc20Abi } from "viem";
import { mainnet } from 'viem/chains'

const publicClient = createPublicClient({
    chain: mainnet,
    transport: http()
});

const overrideBalance = async (balanceSlotIndex : number, walletToOverride : `0x${string}`, newBalanceAmount : bigint, tokenAddress : string) => {
    // Convert the address to a 32-byte hex string
    const paddedAddress = pad(walletToOverride, { size: 32 });
    // Convert the base slot index to a 32-byte hex string
    const paddedBaseSlot = pad(toHex(balanceSlotIndex), { size: 32 });
    // Concatenate the two hex values
    const concatenated = paddedAddress + paddedBaseSlot.slice(2);
    // Calculate the keccak256 hash of the result
    const balanceSlotHash = keccak256(concatenated);
    // New balance: new balance in hex
    const newBalance = pad(numberToHex(newBalanceAmount), { size: 32 });

    const balance = await publicClient.readContract({
        address: tokenAddress,
        abi: erc20Abi,
        functionName: "balanceOf",
        args: [walletToOverride],

        stateOverride: [
            {
                address: tokenAddress,
                stateDiff: [
                    {
                        slot: balanceSlotHash,
                        value: newBalance
                    }
                ]
            }
        ]
    });
};