Change changeBalance matcher type

waffle-chai/src/matchers/changeBalance.ts

@@ -1,9 +1,10 @@
-import {BigNumber, Signer} from 'ethers';
+import {BigNumber} from 'ethers';
+import {Account, getAddressOf, getBalanceOf} from './misc/account';
 
 export function supportChangeBalance(Assertion: Chai.AssertionStatic) {
   Assertion.addMethod('changeBalance', function (
     this: any,
-    signer: Signer,
+    signer: Account,
     balanceChange: any
   ) {
     const subject = this._obj;
@@ -13,13 +14,13 @@ export function supportChangeBalance(Assertion: Chai.AssertionStatic) {
     }
     const derivedPromise = Promise.all([
       getBalanceChange(subject, signer),
-      signer.getAddress()
+      getAddressOf(signer)
     ]).then(
       ([actualChange, address]) => {
         this.assert(
           actualChange.eq(BigNumber.from(balanceChange)),
           `Expected "${address}" to change balance by ${balanceChange} wei, ` +
-            `but it has changed by ${actualChange} wei`,
+          `but it has changed by ${actualChange} wei`,
           `Expected "${address}" to not change balance by ${balanceChange} wei,`,
           balanceChange,
           actualChange
@@ -35,10 +36,11 @@ export function supportChangeBalance(Assertion: Chai.AssertionStatic) {
 
 export async function getBalanceChange(
   transactionCallback: () => any,
-  signer: Signer
+  account: Account
 ) {
-  const balanceBefore = await signer.getBalance();
+  const balanceBefore = await getBalanceOf(account);
   await transactionCallback();
-  const balanceAfter = await signer.getBalance();
+  const balanceAfter = await getBalanceOf(account);
+
   return balanceAfter.sub(balanceBefore);
 }

waffle-chai/src/matchers/changeBalances.ts

@@ -1,9 +1,10 @@
-import {BigNumber, Signer} from 'ethers';
+import {BigNumber} from 'ethers';
+import {getBalanceOf, getAddressOf, Account} from './misc/account';
 
 export function supportChangeBalances(Assertion: Chai.AssertionStatic) {
   Assertion.addMethod('changeBalances', function (
     this: any,
-    signers: Signer[],
+    signers: Account[],
     balanceChanges: any[]
   ) {
     const subject = this._obj;
@@ -38,18 +39,18 @@ export function supportChangeBalances(Assertion: Chai.AssertionStatic) {
 
 async function getBalanceChanges(
   transactionCallback: () => any,
-  signers: Signer[]
+  signers: Account[]
 ) {
   const balancesBefore = await Promise.all(
-    signers.map((signer) => signer.getBalance())
+    signers.map(getBalanceOf)
   );
   await transactionCallback();
   const balancesAfter = await Promise.all(
-    signers.map((signer) => signer.getBalance())
+    signers.map(getBalanceOf)
   );
   return balancesAfter.map((balance, ind) => balance.sub(balancesBefore[ind]));
 }
 
-function getAddresses(signers: Signer[]) {
-  return Promise.all(signers.map((signer) => signer.getAddress()));
+function getAddresses(signers: Account[]) {
+  return Promise.all(signers.map((signer) => getAddressOf(signer)));
 }

waffle-chai/src/matchers/misc/account.ts

@@ -0,0 +1,19 @@
+import {Contract, Signer} from 'ethers';
+
+export type Account = Signer | Contract;
+
+export async function getBalanceOf(account: Account) {
+  if (account instanceof Contract) {
+    return account.provider.getBalance(account.address);
+  } else {
+    return account.getBalance();
+  }
+}
+
+export async function getAddressOf(account: Account) {
+  if (account instanceof Contract) {
+    return account.address;
+  } else {
+    return account.getAddress();
+  }
+}

waffle-chai/test/matchers/balance.test.ts

@@ -1,9 +1,11 @@
 import {expect, AssertionError} from 'chai';
 import {MockProvider} from '@ethereum-waffle/provider';
-import {BigNumber} from 'ethers';
+import {BigNumber, Contract} from 'ethers';
 
 describe('INTEGRATION: Balance observers', () => {
-  const [sender, receiver] = new MockProvider().getWallets();
+  const provider = new MockProvider();
+  const [sender, receiver] = provider.getWallets();
+  const contract = new Contract(receiver.address, [], provider);
 
   describe('Change balance, one account', () => {
     it('Should pass when expected balance change is passed as string and is equal to an actual', async () => {
@@ -159,19 +161,47 @@ describe('INTEGRATION: Balance observers', () => {
     });
 
     it('Should throw when not a callback is passed to expect', async () => {
+      const pendingTx = sender.sendTransaction({
+        to: receiver.address,
+        gasPrice: 0,
+        value: 200
+      });
+
       expect(() =>
         expect(
-          sender.sendTransaction({
-            to: receiver.address,
-            gasPrice: 0,
-            value: 200
-          })
+          pendingTx
         ).to.changeBalances([sender, receiver], ['-200', 200])
       ).to.throw(
         Error,
         `Expect subject should be a callback returning the Promise
         e.g.: await expect(() => wallet.send({to: '0xb', value: 200})).to.changeBalances(['0xa', '0xb'], [-200, 200])`
       );
+
+      await pendingTx;
+    });
+  });
+
+  describe('Change balance, one contract', () => {
+    it('Should pass when expected balance change is passed as int and is equal to an actual', async () => {
+      await expect(async () =>
+        sender.sendTransaction({
+          to: contract.address,
+          gasPrice: 0,
+          value: 200
+        })
+      ).to.changeBalance(contract, 200);
+    });
+  });
+
+  describe('Change balances, one account, one contract', () => {
+    it('Should pass when all expected balance changes are equal to actual values', async () => {
+      await expect(() =>
+        sender.sendTransaction({
+          to: contract.address,
+          gasPrice: 0,
+          value: 200
+        })
+      ).to.changeBalances([sender, contract], [-200, 200]);
     });
   });
 });