Add support for changeTokenBalance(s) matchers

packages/hardhat-chai-matchers/src/changeTokenBalance.ts

@@ -0,0 +1,177 @@
+import { BigNumber, BigNumberish, Contract, providers } from "ethers";
+import { ensure } from "./calledOnContract/utils";
+import { Account, getAddressOf } from "./misc/account";
+
+type TransactionResponse = providers.TransactionResponse;
+
+interface Token extends Contract {
+  balanceOf(address: string, overrides?: any): Promise<BigNumber>;
+}
+
+export function supportChangeTokenBalance(Assertion: Chai.AssertionStatic) {
+  Assertion.addMethod(
+    "changeTokenBalance",
+    function (
+      this: any,
+      token: Token,
+      account: Account | string,
+      balanceChange: BigNumberish
+    ) {
+      const subject = this._obj;
+
+      checkToken(token, "changeTokenBalance");
+
+      const derivedPromise = Promise.all([
+        getBalanceChange(subject, token, account),
+        getAddressOf(account),
+        getTokenDescription(token),
+      ]).then(([actualChange, address, tokenDescription]) => {
+        this.assert(
+          actualChange.eq(BigNumber.from(balanceChange)),
+          `Expected "${address}" to change its balance of ${tokenDescription} by ${balanceChange}, ` +
+            `but it has changed by ${actualChange}`,
+          `Expected "${address}" to not change its balance of ${tokenDescription} by ${balanceChange}, but it did`,
+          balanceChange,
+          actualChange
+        );
+      });
+
+      this.then = derivedPromise.then.bind(derivedPromise);
+      this.catch = derivedPromise.catch.bind(derivedPromise);
+
+      return this;
+    }
+  );
+
+  Assertion.addMethod(
+    "changeTokenBalances",
+    function (
+      this: any,
+      token: Token,
+      accounts: Array<Account | string>,
+      balanceChanges: BigNumberish[]
+    ) {
+      const subject = this._obj;
+
+      checkToken(token, "changeTokenBalances");
+
+      if (accounts.length !== balanceChanges.length) {
+        throw new Error(
+          `The number of accounts (${accounts.length}) is different than the number of expected balance changes (${balanceChanges.length})`
+        );
+      }
+
+      const balanceChangesPromise = Promise.all(
+        accounts.map((account) => getBalanceChange(subject, token, account))
+      );
+      const addressesPromise = Promise.all(accounts.map(getAddressOf));
+
+      const derivedPromise = Promise.all([
+        balanceChangesPromise,
+        addressesPromise,
+        getTokenDescription(token),
+      ]).then(([actualChanges, addresses, tokenDescription]) => {
+        this.assert(
+          actualChanges.every((change, ind) =>
+            change.eq(BigNumber.from(balanceChanges[ind]))
+          ),
+          `Expected ${addresses} to change their balance of ${tokenDescription} by ${balanceChanges}, ` +
+            `but it has changed by ${actualChanges}`,
+          `Expected ${addresses} to not change their balance of ${tokenDescription} by ${balanceChanges}, but they did`,
+          balanceChanges.map((balanceChange) => balanceChange.toString()),
+          actualChanges.map((actualChange) => actualChange.toString())
+        );
+      });
+
+      this.then = derivedPromise.then.bind(derivedPromise);
+      this.catch = derivedPromise.catch.bind(derivedPromise);
+
+      return this;
+    }
+  );
+}
+
+function checkToken(token: unknown, method: string) {
+  if (typeof token !== "object" || token === null || !("functions" in token)) {
+    throw new Error(
+      `The first argument of ${method} must be the contract instance of the token`
+    );
+  } else if ((token as any).functions.balanceOf === undefined) {
+    throw new Error("The given contract instance is not an ERC20 token");
+  }
+}
+
+export async function getBalanceChange(
+  transaction:
+    | TransactionResponse
+    | Promise<TransactionResponse>
+    | (() => TransactionResponse)
+    | (() => Promise<TransactionResponse>),
+  token: Token,
+  account: Account | string
+) {
+  const hre = await import("hardhat");
+  const provider = hre.network.provider;
+
+  let txResponse: TransactionResponse;
+
+  if (typeof transaction === "function") {
+    txResponse = await transaction();
+  } else {
+    txResponse = await transaction;
+  }
+
+  const txReceipt = await txResponse.wait();
+  const txBlockNumber = txReceipt.blockNumber;
+
+  const block = await provider.send("eth_getBlockByHash", [
+    txReceipt.blockHash,
+    false,
+  ]);
+
+  ensure(
+    block.transactions.length === 1,
+    Error,
+    "Multiple transactions found in block"
+  );
+
+  const address = await getAddressOf(account);
+
+  const balanceAfter = await token.balanceOf(address, {
+    blockTag: txBlockNumber,
+  });
+
+  const balanceBefore = await token.balanceOf(address, {
+    blockTag: txBlockNumber - 1,
+  });
+
+  return BigNumber.from(balanceAfter).sub(balanceBefore);
+}
+
+let tokenDescriptionsCache: Record<string, string> = {};
+/**
+ * Get a description for the given token. Use the symbol of the token if
+ * possible; if it doesn't exist, the name is used; if the name doesn't
+ * exist, the address of the token is used.
+ */
+async function getTokenDescription(token: Token): Promise<string> {
+  if (tokenDescriptionsCache[token.address] === undefined) {
+    let tokenDescription = `<token at ${token.address}>`;
+    try {
+      tokenDescription = await token.symbol();
+    } catch (e) {
+      try {
+        tokenDescription = await token.name();
+      } catch (e2) {}
+    }
+
+    tokenDescriptionsCache[token.address] = tokenDescription;
+  }
+
+  return tokenDescriptionsCache[token.address];
+}
+
+// only used by tests
+export function clearTokenDescriptionsCache() {
+  tokenDescriptionsCache = {};
+}

packages/hardhat-chai-matchers/src/hardhatChaiMatchers.ts

@@ -6,6 +6,7 @@ import { supportProperAddress } from "./properAddress";
 import { supportProperPrivateKey } from "./properPrivateKey";
 import { supportChangeEtherBalance } from "./changeEtherBalance";
 import { supportChangeEtherBalances } from "./changeEtherBalances";
+import { supportChangeTokenBalance } from "./changeTokenBalance";
 import { supportReverted } from "./reverted/reverted";
 import { supportRevertedWith } from "./reverted/revertedWith";
 import { supportRevertedWithCustomError } from "./reverted/revertedWithCustomError";
@@ -23,6 +24,7 @@ export function hardhatChaiMatchers(
   supportProperPrivateKey(chai.Assertion);
   supportChangeEtherBalance(chai.Assertion);
   supportChangeEtherBalances(chai.Assertion);
+  supportChangeTokenBalance(chai.Assertion);
   supportReverted(chai.Assertion);
   supportRevertedWith(chai.Assertion);
   supportRevertedWithCustomError(chai.Assertion, utils);

packages/hardhat-chai-matchers/src/types.ts

@@ -26,6 +26,12 @@ declare namespace Chai {
       balances: any[],
       options?: any
     ): AsyncAssertion;
+    changeTokenBalance(token: any, account: any, balance: any): AsyncAssertion;
+    changeTokenBalances(
+      token: any,
+      account: any[],
+      balance: any[]
+    ): AsyncAssertion;
   }
 
   interface NumericComparison {