From e03b3a6afeb85e39286c831ee67f58953d68da45 Mon Sep 17 00:00:00 2001
From: Paul Razvan Berg <hello@paulrberg.com>
Date: Mon, 24 Aug 2020 20:10:25 +0300
Subject: [PATCH] feat: revert with reason in doppelganger contract

---
 waffle-mock-contract/src/Doppelganger.sol  | 11 +++++++----
 waffle-mock-contract/src/index.ts          |  3 ++-
 waffle-mock-contract/test/contract.test.ts | 11 +++++++++--
 3 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/waffle-mock-contract/src/Doppelganger.sol b/waffle-mock-contract/src/Doppelganger.sol
index 7f2f705d7..066637fdf 100644
--- a/waffle-mock-contract/src/Doppelganger.sol
+++ b/waffle-mock-contract/src/Doppelganger.sol
@@ -5,6 +5,7 @@ contract Doppelganger {
     struct MockCall {
         bool initialized;
         bool reverts;
+        string revertReason;
         bytes returnValue;
     }
 
@@ -13,16 +14,17 @@ contract Doppelganger {
     fallback() external payable {
         MockCall storage mockCall = __internal__getMockCall();
         if (mockCall.reverts == true) {
-            __internal__mockRevert();
+            __internal__mockRevert(mockCall.revertReason);
             return;
         }
         __internal__mockReturn(mockCall.returnValue);
     }
 
-    function __waffle__mockReverts(bytes memory data) public {
+    function __waffle__mockReverts(bytes memory data, string memory reason) public {
         mockConfig[keccak256(data)] = MockCall({
             initialized: true,
             reverts: true,
+            revertReason: reason,
             returnValue: ""
         });
     }
@@ -31,6 +33,7 @@ contract Doppelganger {
         mockConfig[keccak256(data)] = MockCall({
             initialized: true,
             reverts: false,
+            revertReason: "",
             returnValue: value
         });
     }
@@ -67,7 +70,7 @@ contract Doppelganger {
         }
     }
 
-    function __internal__mockRevert() pure private {
-        revert("Mock revert");
+    function __internal__mockRevert(string memory reason) pure private {
+        revert(reason);
     }
 }
diff --git a/waffle-mock-contract/src/index.ts b/waffle-mock-contract/src/index.ts
index 4abac4ce1..9bc6eefc1 100644
--- a/waffle-mock-contract/src/index.ts
+++ b/waffle-mock-contract/src/index.ts
@@ -21,7 +21,8 @@ function stub(mockContract: Contract, encoder: utils.AbiCoder, func: utils.Funct
       const encoded = encoder.encode(func.outputs, args);
       await mockContract.__waffle__mockReturns(callData, encoded);
     },
-    reverts: async () => mockContract.__waffle__mockReverts(callData),
+    reverts: async () => mockContract.__waffle__mockReverts(callData, 'Mock revert'),
+    revertsWithReason: async (reason: string) => mockContract.__waffle__mockReverts(callData, reason),
     withArgs: (...args: any[]) => stub(mockContract, encoder, func, args)
   };
 }
diff --git a/waffle-mock-contract/test/contract.test.ts b/waffle-mock-contract/test/contract.test.ts
index 05ed9ef18..5f47f1388 100644
--- a/waffle-mock-contract/test/contract.test.ts
+++ b/waffle-mock-contract/test/contract.test.ts
@@ -105,10 +105,17 @@ describe('Doppelganger - Contract', () => {
       expect(await pretender.testArgumentTypes(1, false, 'str', '0x0123')).to.equal('0x0123');
     });
 
+    it('reverts with no message', async () => {
+      const {contract, pretender} = await deploy();
+
+      await contract.__waffle__mockReverts(readSignature, '');
+      await expect(pretender.read()).to.be.revertedWith('');
+    });
+
     it('reverts with correct message', async () => {
       const {contract, pretender} = await deploy();
 
-      await contract.__waffle__mockReverts(readSignature);
+      await contract.__waffle__mockReverts(readSignature, 'Mock revert');
       await expect(pretender.read()).to.be.revertedWith('Mock revert');
     });
 
@@ -118,7 +125,7 @@ describe('Doppelganger - Contract', () => {
       const callData = `${addSignature}0000000000000000000000000000000000000000000000000000000000000005`;
       const returnedValue = '0x1000000000000000000000000000000000000000000000000000000000004234';
 
-      await contract.__waffle__mockReverts(callData);
+      await contract.__waffle__mockReverts(callData, 'Mock revert');
       await contract.__waffle__mockReturns(addSignature, returnedValue);
 
       await expect(pretender.add(5)).to.be.revertedWith('Mock revert');