🍶 Implement mocking receive function to revert (#807)

.changeset/quiet-bugs-jam.md

@@ -0,0 +1,5 @@
+---
+"@ethereum-waffle/mock-contract": patch
+---
+
+🍶 Implement mocking receive function to revert

docs/source/mock-contract.rst

@@ -122,3 +122,67 @@ Mock contract will be used to mock exactly this call with values that are releva
       expect(await contract.connect(receiver.address).check()).to.equal(false);
     });
   });
+
+Mocking receive function
+------------------------
+
+The :code:`receive` function of the mocked Smart Contract can be mocked to revert. It cannot however be mocked to return a specified value, because of gas limitations when calling another contract using :code:`send` and :code:`transfer`.
+
+Receive mock example
+^^^^^^^^^^^^^^^^^^^^
+
+.. code-block:: solidity
+
+  pragma solidity ^0.6.0;
+
+  interface IERC20 {
+      function balanceOf(address account) external view returns (uint256);
+      fallback() external payable;
+      receive() external payable;
+  }
+
+  contract EtherForward {
+      IERC20 private tokenContract;
+
+      constructor (IERC20 _tokenContract) public {
+          tokenContract = _tokenContract;
+      }
+
+      function forward() public payable {
+          payable(tokenContract).transfer(msg.value);
+      }
+  }
+
+.. code-block:: ts
+
+  (...)
+
+  it('use the receive function normally', async () => {
+    const {contract, mockERC20} = await setup();
+
+    expect (
+      await mockERC20.provider.getBalance(mockERC20.address)
+    ).to.be.equal(0);
+
+    await contract.forward({value: 7})
+
+    expect (
+      await mockERC20.provider.getBalance(mockERC20.address)
+    ).to.be.equal(7);
+  });
+
+  it('can mock the receive function to revert', async () => {
+    const {contract, mockERC20} = await setup();
+
+    await mockERC20.mock.receive.revertsWithReason('Receive function rejected')
+
+    await expect(
+      contract.forward({value: 7})
+    ).to.be.revertedWith('Receive function rejected')
+
+    expect (
+      await mockERC20.provider.getBalance(mockERC20.address)
+    ).to.be.equal(0);
+  });
+
+  (...)

waffle-mock-contract/src/Doppelganger.sol

@@ -10,6 +10,8 @@ contract Doppelganger {
     }
 
     mapping(bytes32 => MockCall) mockConfig;
+    bool receiveReverts;
+    string receiveRevertReason;
     mapping(bytes32 => bytes32) heads;
 
     fallback() external payable {
@@ -21,6 +23,10 @@ contract Doppelganger {
         __internal__mockReturn(mockCall.returnValue);
     }
 
+    receive() payable external {
+        require(receiveReverts == false, receiveRevertReason);
+    }
+
     function __clearQueue(bytes32 at) private {
         heads[at] = at;
         while(mockConfig[at].next != "") {
@@ -66,6 +72,16 @@ contract Doppelganger {
         __waffle__queueReturn(data, value);
     }
 
+    function __waffle__receiveReverts(string memory reason) public {
+        receiveReverts = true;
+        receiveRevertReason = reason;
+    }
+
+    function __waffle__receiveReverts(string memory reason) public {
+        receiveReverts = true;
+        receiveRevertReason = reason;
+    }
+
     function __waffle__call(address target, bytes calldata data) external returns (bytes memory) {
       (bool succeeded, bytes memory returnValue) = target.call(data);
       require(succeeded, string(returnValue));

waffle-mock-contract/src/index.ts

@@ -163,6 +163,13 @@ function createMock(abi: ABI, mockContractInstance: Contract) {
     };
   }, {} as MockContract['mock']);
 
+  mockedAbi.receive = {
+    returns: async () => { throw new Error('Receive function return is not implemented.'); },
+    withArgs: () => { throw new Error('Receive function return is not implemented.'); },
+    reverts: async () => mockContractInstance.__waffle__receiveReverts('Mock Revert'),
+    revertsWithReason: async (reason: string) => mockContractInstance.__waffle__receiveReverts(reason)
+  };
+
   return mockedAbi;
 }
 

waffle-mock-contract/test/amirichalready.test.ts

@@ -2,7 +2,7 @@ import {use, expect} from 'chai';
 import {Contract, ContractFactory, utils, Wallet} from 'ethers';
 import {MockProvider} from '@ethereum-waffle/provider';
 import {waffleChai} from '@ethereum-waffle/chai';
-import {deployMockContract} from '../src';
+import {deployMockContract, MockContract} from '../src';
 
 import IERC20 from './helpers/interfaces/IERC20.json';
 import AmIRichAlready from './helpers/interfaces/AmIRichAlready.json';
@@ -13,7 +13,7 @@ describe('Am I Rich Already', () => {
   let contractFactory: ContractFactory;
   let sender: Wallet;
   let receiver: Wallet;
-  let mockERC20: Contract;
+  let mockERC20: MockContract;
   let contract: Contract;
 
   beforeEach(async () => {

waffle-mock-contract/test/etherForward.test.ts

@@ -0,0 +1,64 @@
+import {waffleChai} from '@ethereum-waffle/chai';
+import {MockProvider} from '@ethereum-waffle/provider';
+import {expect, use} from 'chai';
+import {Contract, ContractFactory, Wallet} from 'ethers';
+import {deployMockContract} from '../src';
+
+import EtherForward from './helpers/interfaces/EtherForward.json';
+import IERC20 from './helpers/interfaces/IERC20.json';
+
+use(waffleChai);
+
+describe('Ether Forwarded', () => {
+  let contractFactory: ContractFactory;
+  let sender: Wallet;
+  let mockERC20: Contract;
+  let contract: Contract;
+  let provider: MockProvider;
+
+  beforeEach(async () => {
+    provider = new MockProvider()
+    ;[sender] = provider.getWallets();
+    mockERC20 = await deployMockContract(sender, IERC20.abi);
+    contractFactory = new ContractFactory(EtherForward.abi, EtherForward.bytecode, sender);
+    contract = await contractFactory.deploy(mockERC20.address);
+  });
+
+  it('Can forward ether through call', async () => {
+    expect(await provider.getBalance(mockERC20.address)).to.be.equal(0);
+    await contract.forwardByCall({value: 7});
+    expect(await provider.getBalance(mockERC20.address)).to.be.equal(7);
+  });
+
+  it('Can forward ether through send', async () => {
+    expect(await provider.getBalance(mockERC20.address)).to.be.equal(0);
+    await contract.forwardBySend({value: 7});
+    expect(await provider.getBalance(mockERC20.address)).to.be.equal(7);
+  });
+
+  it('Can forward ether through transfer', async () => {
+    expect(await provider.getBalance(mockERC20.address)).to.be.equal(0);
+    await contract.forwardByTransfer({value: 7});
+    expect(await provider.getBalance(mockERC20.address)).to.be.equal(7);
+  });
+
+  it('Can mock a revert on a receive function', async () => {
+    expect(await provider.getBalance(mockERC20.address)).to.be.equal(0);
+
+    await mockERC20.mock.receive.revertsWithReason('Receive function rejected ether.');
+
+    await expect(
+      contract.forwardByCall({value: 7})
+    ).to.be.revertedWith('Receive function rejected ether.');
+
+    await expect(
+      contract.forwardBySend({value: 7})
+    ).to.be.revertedWith('forwardBySend failed');
+
+    await expect(
+      contract.forwardByTransfer({value: 7})
+    ).to.be.reverted;
+
+    expect(await provider.getBalance(mockERC20.address)).to.be.equal(0);
+  });
+});

waffle-mock-contract/test/helpers/contracts/EtherForward.sol

@@ -0,0 +1,34 @@
+pragma solidity ^0.6.3;
+
+interface IERC20 {
+    function balanceOf(address account) external view returns (uint256);
+    fallback() external payable;
+    receive() external payable;
+}
+
+contract EtherForward {
+    IERC20 private tokenContract;
+
+    constructor (IERC20 _tokenContract) public {
+        tokenContract = _tokenContract;
+    }
+
+    function forwardByCall() public payable {
+        (bool sent, bytes memory data) = payable(tokenContract).call{value: msg.value}("");
+        if (!sent) {
+            // https://ethereum.stackexchange.com/a/114140/24330
+            // Bubble up the revert from the call.
+            assembly {
+                revert(add(data, 32), data)
+            }
+        }
+    }
+
+    function forwardBySend() public payable {
+        require(payable(tokenContract).send(msg.value), "forwardBySend failed");
+    }
+
+    function forwardByTransfer() public payable {
+        payable(tokenContract).transfer(msg.value);
+    }
+}