Repository: defi-wonderland/smock (ts)

Issue:
EVM Reset bug
**Describe the bug**
Calling `evm.reset()` blocks the creation of a new fake.

**Reproduction steps**

```typescript
const fake1 = await smock.fake('MyContract');
await fake1.callStatic.something(); // will be fine

await network.provider.request({
    method: 'hardhat_reset',
    params: [],
});

const fake2 = await smock.fake('MyOtherContract');
await fake2.callStatic.somethingElse(); // will revert
```

**Expected behavior**
`fake2` should be mocked, and behave as such. N items should be able to be mocked.

**System Specs:**

- OS: Ubuntu 20.04
- Package Version (or commit hash): This happend with hardhat 2.5.0, and also with 2.6.1. Also ethers 5.4.1 and 5.4.3

**Additional context**
Showed with @0xGorilla how to reproduce.


Pull request:
fix: reinit sandbox is case of an evm reset
**Description**
Closes #55 


Patch:
diff --git a/docs/source/fakes.rst b/docs/source/fakes.rst
index 8d9a2c9..3f7d138 100644
--- a/docs/source/fakes.rst
+++ b/docs/source/fakes.rst
@@ -89,6 +89,16 @@ Options
   }
 
 
+Signing transactions
+********************
+
+Every fake comes with a :code:`wallet` property in order to make easy to sign transactions
+
+.. code-block:: typescript
+
+  myContract.connect(myFake.wallet).doSomething();
+
+
 Making a function return
 ************************
 
@@ -139,6 +149,15 @@ Returning a dynamic value
 Returning a value based on arguments
 ####################################
 
+.. code-block:: typescript
+
+  myFake.myFunction.whenCalledWith(123).returns(456);
+  
+  await myFake.myFunction(123); // returns 456
+
+Returning a value with custom logic
+###################################
+
 .. code-block:: typescript
 
   myFake.getDynamicInput.returns(arg1 => arg1 * 10);
@@ -192,6 +211,18 @@ Reverting at a specific call count
   await myFake.myFunction(); // reverts with 'Something went wrong'
   await myFake.myFunction(); // returns 1234
 
+Reverting based on arguments
+############################
+
+.. code-block:: typescript
+
+  myFake.myFunction.returns(1);
+  myFake.myFunction.whenCalledWith(123).reverts('Something went wrong');
+
+  await myFake.myFunction(); // returns 1
+  await myFake.myFunction(123); // reverts with 'Something went wrong'
+
+
 Resetting function behavior
 ***************************
 
@@ -328,14 +359,14 @@ Manipulating fallback functions
 *******************************
 
 Modifying the :code:`fallback` function
-#################################
+#######################################
 
 .. code-block:: typescript
 
   myFake.fallback.returns();
 
 Modifying the :code:`receive` function
-################################
+######################################
 
 .. code-block:: typescript
 
diff --git a/docs/source/mocks.rst b/docs/source/mocks.rst
index 28a0b07..3c978bb 100644
--- a/docs/source/mocks.rst
+++ b/docs/source/mocks.rst
@@ -94,3 +94,25 @@ Setting the value of a struct
     valueA: 1234,
     valueB: true,
   });
+
+
+Setting the value of a mapping (won't affect other keys)
+#######################################################
+
+.. code-block:: typescript
+
+  await myMock.setVariable('myMapping', {
+    myKey: 1234
+  });
+
+
+Setting the value of a nested mapping
+#####################################
+
+.. code-block:: typescript
+
+  await myMock.setVariable('myMapping', {
+    myChildMapping: {
+      myKey: 1234
+    }
+  });
\ No newline at end of file
diff --git a/src/factories/smock-contract.ts b/src/factories/smock-contract.ts
index f0ec287..3db8879 100644
--- a/src/factories/smock-contract.ts
+++ b/src/factories/smock-contract.ts
@@ -140,6 +140,7 @@ function fillProgrammableContractFunction(fn: ProgrammableContractFunction, logi
   fn.returnsAtCall = logic.returnsAtCall.bind(logic);
   fn.reverts = logic.reverts.bind(logic);
   fn.revertsAtCall = logic.revertsAtCall.bind(logic);
+  fn.whenCalledWith = logic.whenCalledWith.bind(logic);
   fn.reset = logic.reset.bind(logic);
 }
 
diff --git a/src/index.ts b/src/index.ts
index 46406e5..74a6037 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -1,19 +1,31 @@
 import { BaseContract, ContractFactory } from 'ethers';
+import hre from 'hardhat';
 import { matchers } from './chai-plugin/matchers';
 import { Sandbox } from './sandbox';
 import { FakeContract, FakeContractOptions, FakeContractSpec, MockContractFactory } from './types';
+import { getHardhatBaseProvider } from './utils';
 
 let sandbox: Sandbox;
 
 async function fake<T extends BaseContract>(spec: FakeContractSpec, opts: FakeContractOptions = {}): Promise<FakeContract<T>> {
-  if (!sandbox) sandbox = await Sandbox.create();
+  if (!sandbox) await init();
   return await sandbox.fake(spec, opts);
 }
 
 async function mock<T extends ContractFactory>(contractName: string): Promise<MockContractFactory<T>> {
-  if (!sandbox) sandbox = await Sandbox.create();
+  if (!sandbox) await init();
+
   return await sandbox.mock(contractName);
 }
 
+async function init() {
+  sandbox = await Sandbox.create();
+
+  // reinit in case there is an evm reset
+  (await getHardhatBaseProvider(hre)).on('hardhatNetworkReset', async () => {
+    sandbox = await Sandbox.create();
+  });
+}
+
 export * from './types';
 export const smock = { fake, mock, matchers };
diff --git a/src/logic/programmable-function-logic.ts b/src/logic/programmable-function-logic.ts
index f3d1714..7e0e832 100644
--- a/src/logic/programmable-function-logic.ts
+++ b/src/logic/programmable-function-logic.ts
@@ -2,8 +2,9 @@ import { EVMResult } from '@nomiclabs/ethereumjs-vm/dist/evm/evm';
 import { VmError } from '@nomiclabs/ethereumjs-vm/dist/exceptions';
 import BN from 'bn.js';
 import { ethers } from 'ethers';
+import { findLast } from 'lodash';
 import { Observable, withLatestFrom } from 'rxjs';
-import { ContractCall, ProgrammedReturnValue } from '../index';
+import { ContractCall, ProgrammedReturnValue, WhenCalledWithChain } from '../index';
 import { WatchableFunctionLogic } from '../logic/watchable-function-logic';
 import { fromHexString } from '../utils';
 
@@ -23,6 +24,7 @@ export class ProgrammableFunctionLogic extends WatchableFunctionLogic {
   protected encoder: (value: any) => string;
   protected defaultAnswer: ProgrammedAnswer | undefined;
   protected answerByIndex: { [index: number]: ProgrammedAnswer } = {};
+  protected answerByArgs: { answer: ProgrammedAnswer; args: unknown[] }[] = [];
 
   constructor(
     name: string,
@@ -49,6 +51,23 @@ export class ProgrammableFunctionLogic extends WatchableFunctionLogic {
     this.answerByIndex[callIndex] = new ProgrammedAnswer(value, false);
   }
 
+  whenCalledWith(...args: unknown[]): WhenCalledWithChain {
+    return {
+      returns: (value?: ProgrammedReturnValue) => {
+        this.answerByArgs.push({
+          args,
+          answer: new ProgrammedAnswer(value, false),
+        });
+      },
+      reverts: (reason?: string) => {
+        this.answerByArgs.push({
+          args,
+          answer: new ProgrammedAnswer(reason, true),
+        });
+      },
+    };
+  }
+
   reverts(reason?: string): void {
     this.defaultAnswer = new ProgrammedAnswer(reason, true);
   }
@@ -61,10 +80,11 @@ export class ProgrammableFunctionLogic extends WatchableFunctionLogic {
     super.reset();
     this.defaultAnswer = undefined;
     this.answerByIndex = {};
+    this.answerByArgs = [];
   }
 
   private async modifyAnswer(result: EVMResult, call: ContractCall): Promise<void> {
-    const answer = this.answerByIndex[this.getCallCount() - 1] || this.defaultAnswer;
+    const answer = this.getCallAnswer(call);
 
     if (answer) {
       result.gasUsed = new BN(0);
@@ -78,6 +98,21 @@ export class ProgrammableFunctionLogic extends WatchableFunctionLogic {
     }
   }
 
+  private getCallAnswer(call: ContractCall): ProgrammedAnswer | undefined {
+    let answer: ProgrammedAnswer | undefined;
+
+    // if there is an answer for this call index, return it
+    answer = this.answerByIndex[this.getCallCount() - 1];
+    if (answer) return answer;
+
+    // if there is an answer for this call arguments, return it
+    answer = findLast(this.answerByArgs, (option) => this.isDeepEqual(option.args, call.args))?.answer;
+    if (answer) return answer;
+
+    // return the default answer
+    return this.defaultAnswer;
+  }
+
   private async encodeValue(value: ProgrammedReturnValue, call: ContractCall): Promise<Buffer> {
     if (value === undefined) return EMPTY_ANSWER;
 
diff --git a/src/logic/watchable-function-logic.ts b/src/logic/watchable-function-logic.ts
index 7e0496b..33f0591 100644
--- a/src/logic/watchable-function-logic.ts
+++ b/src/logic/watchable-function-logic.ts
@@ -28,11 +28,11 @@ export class WatchableFunctionLogic {
   }
 
   calledWith(...expectedCallArgs: unknown[]): boolean {
-    return !!this.callHistory.find((call) => isEqualWith(call.args, expectedCallArgs, this.isEqualCustomizer.bind(this)));
+    return !!this.callHistory.find((call) => this.isDeepEqual(call.args, expectedCallArgs));
   }
 
   alwaysCalledWith(...expectedCallArgs: unknown[]): boolean {
-    const callWithOtherArgs = this.callHistory.find((call) => !isEqualWith(call.args, expectedCallArgs, this.isEqualCustomizer.bind(this)));
+    const callWithOtherArgs = this.callHistory.find((call) => !this.isDeepEqual(call.args, expectedCallArgs));
     return this.getCalled() && !callWithOtherArgs;
   }
 
@@ -134,6 +134,10 @@ export class WatchableFunctionLogic {
     });
   }
 
+  protected isDeepEqual(obj1: unknown, obj2: unknown): boolean {
+    return isEqualWith(obj1, obj2, this.isEqualCustomizer.bind(this));
+  }
+
   /**
    * Normal deep comparison wont be enought when comparing expect vs result objects
    * - In case we are comparing 2 big numbers, we must compare them using the BigNumber methods
diff --git a/src/types.ts b/src/types.ts
index 8609c71..6cd4d4e 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -34,11 +34,17 @@ export interface ContractCall {
   target: string;
 }
 
+export type WhenCalledWithChain = {
+  returns: (value?: ProgrammedReturnValue) => void;
+  reverts: (reason?: string) => void;
+};
+
 export interface ProgrammableContractFunction extends WatchableContractFunction {
   returns: (value?: ProgrammedReturnValue) => void;
   returnsAtCall: (index: number, value?: ProgrammedReturnValue) => void;
   reverts: (reason?: string) => void;
   revertsAtCall: (index: number, reason?: string) => void;
+  whenCalledWith: (...args: unknown[]) => WhenCalledWithChain;
   reset: () => void;
 }
 


Test patch:
diff --git a/test/contracts/programmable-function-logic/PickyReturner.sol b/test/contracts/programmable-function-logic/PickyReturner.sol
new file mode 100644
index 0000000..4f4bb8c
--- /dev/null
+++ b/test/contracts/programmable-function-logic/PickyReturner.sol
@@ -0,0 +1,6 @@
+// SPDX-License-Identifier: MIT
+pragma solidity ^0.8.4;
+
+contract PickyReturner {
+  function getUint256(uint256) public returns (uint256) {}
+}
diff --git a/test/unit/fake/initialization.spec.ts b/test/unit/fake/initialization.spec.ts
index 7181e67..a32d45c 100644
--- a/test/unit/fake/initialization.spec.ts
+++ b/test/unit/fake/initialization.spec.ts
@@ -3,7 +3,7 @@ import { makeRandomAddress } from '@src/utils';
 import { Receiver, Receiver__factory, Returner } from '@typechained';
 import receiverArtifact from 'artifacts/test/contracts/watchable-function-logic/Receiver.sol/Receiver.json';
 import chai, { expect } from 'chai';
-import { ethers } from 'hardhat';
+import { ethers, network } from 'hardhat';
 
 chai.use(smock.matchers);
 
@@ -63,4 +63,17 @@ describe('Fake: Initialization', () => {
     const returner = (await ethers.getContractAt('Returner', targetAddress)) as Returner;
     expect(await returner.callStatic.getBoolean()).to.be.true;
   });
+
+  it('should handle evm resets', async () => {
+    const receiver1 = await smock.fake<Receiver>('Receiver');
+    await receiver1.callStatic.receiveEmpty();
+
+    await network.provider.request({
+      method: 'hardhat_reset',
+      params: [],
+    });
+
+    const receiver2 = await smock.fake<Returner>('Returner');
+    await expect(receiver2.callStatic.getBoolean()).not.to.be.reverted;
+  });
 });
diff --git a/test/unit/programmable-function-logic/when-called-with.spec.ts b/test/unit/programmable-function-logic/when-called-with.spec.ts
new file mode 100644
index 0000000..7810807
--- /dev/null
+++ b/test/unit/programmable-function-logic/when-called-with.spec.ts
@@ -0,0 +1,90 @@
+import { FakeContract, smock } from '@src';
+import { PickyReturner } from '@typechained';
+import { expect } from 'chai';
+
+describe('ProgrammableFunctionLogic: When called with', () => {
+  let fake: FakeContract<PickyReturner>;
+
+  beforeEach(async () => {
+    fake = await smock.fake<PickyReturner>('PickyReturner');
+  });
+
+  describe('returns', () => {
+    it('should override default behaviour', async () => {
+      fake.getUint256.returns(1);
+      fake.getUint256.whenCalledWith(123).returns(456);
+
+      expect(await fake.callStatic.getUint256(123)).to.equal(456);
+    });
+
+    it('should override itself', async () => {
+      fake.getUint256.whenCalledWith(123).returns(1);
+      fake.getUint256.whenCalledWith(123).returns(2);
+
+      expect(await fake.callStatic.getUint256(123)).to.equal(2);
+    });
+
+    it('should be reseted', async () => {
+      fake.getUint256.whenCalledWith(123).returns(1);
+      fake.getUint256.reset();
+      fake.getUint256.returns(2);
+
+      expect(await fake.callStatic.getUint256(123)).to.equal(2);
+    });
+
+    it('should live together with default', async () => {
+      fake.getUint256.returns(1);
+      fake.getUint256.whenCalledWith(123).returns(456);
+
+      expect(await fake.callStatic.getUint256(122)).to.equal(1);
+      expect(await fake.callStatic.getUint256(123)).to.equal(456);
+    });
+
+    it('should handle multiple calls', async () => {
+      fake.getUint256.whenCalledWith(1).returns(10);
+      fake.getUint256.whenCalledWith(2).returns(20);
+
+      expect(await fake.callStatic.getUint256(1)).to.equal(10);
+      expect(await fake.callStatic.getUint256(2)).to.equal(20);
+    });
+  });
+
+  describe('reverts', () => {
+    it('should override default behaviour', async () => {
+      fake.getUint256.reverts('a');
+      fake.getUint256.whenCalledWith(123).reverts('b');
+
+      await expect(fake.callStatic.getUint256(123)).to.be.revertedWith('b');
+    });
+
+    it('should override itself', async () => {
+      fake.getUint256.whenCalledWith(123).reverts('a');
+      fake.getUint256.whenCalledWith(123).reverts('b');
+
+      await expect(fake.callStatic.getUint256(123)).to.be.revertedWith('b');
+    });
+
+    it('should be reseted', async () => {
+      fake.getUint256.whenCalledWith(123).reverts('a');
+      fake.getUint256.reset();
+
+      await expect(fake.callStatic.getUint256(123)).not.to.be.reverted;
+    });
+
+    it('should live together with default', async () => {
+      fake.getUint256.reverts('a');
+      fake.getUint256.whenCalledWith(123).reverts('b');
+
+      await expect(fake.callStatic.getUint256(122)).to.be.revertedWith('a');
+      await expect(fake.callStatic.getUint256(123)).to.be.revertedWith('b');
+    });
+
+    it('should handle multiple calls', async () => {
+      fake.getUint256.whenCalledWith(1).reverts('a');
+      fake.getUint256.whenCalledWith(2).reverts('b');
+
+      await expect(fake.callStatic.getUint256(1)).to.be.revertedWith('a');
+      await expect(fake.callStatic.getUint256(2)).to.be.revertedWith('b');
+    });
+  });
+});