Repository: defi-wonderland/smock (ts)

Issue:
connect function leaves types behind
**Describe the bug**
When doing `myContract.connect(someone)`, the type is lost, leaving a `BaseContract`

**Expected behavior**
`connect` should return a `FakeContract` or a `MockContract`

Pull request description:
feat: propagate types when using connect function
**Description**
Closes #60 


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..651a1c9 100644
--- a/src/factories/smock-contract.ts
+++ b/src/factories/smock-contract.ts
@@ -18,7 +18,7 @@ export async function createFakeContract<Contract extends BaseContract>(
   contractInterface: ethers.utils.Interface,
   provider: ethers.providers.Provider
 ): Promise<FakeContract<Contract>> {
-  const fake = await initContract<FakeContract<Contract>>(vm, address, contractInterface, provider);
+  const fake = (await initContract(vm, address, contractInterface, provider)) as unknown as FakeContract<Contract>;
   const contractFunctions = getContractFunctionsNameAndSighash(contractInterface, Object.keys(fake.functions));
 
   // attach to every contract function, all the programmable and watchable logic
@@ -35,7 +35,7 @@ export async function createMockContractFactory<T extends ContractFactory>(
   vm: ObservableVM,
   contractName: string
 ): Promise<MockContractFactory<T>> {
-  const factory = (await hardhatEthers.getContractFactory(contractName)) as MockContractFactory<T>;
+  const factory = (await hardhatEthers.getContractFactory(contractName)) as unknown as MockContractFactory<T>;
 
   const realDeploy = factory.deploy;
   factory.deploy = async (...args: Parameters<T['deploy']>) => {
@@ -59,15 +59,15 @@ export async function createMockContractFactory<T extends ContractFactory>(
   return factory;
 }
 
-async function initContract<T extends BaseContract>(
+async function initContract(
   vm: ObservableVM,
   address: string,
   contractInterface: ethers.utils.Interface,
   provider: ethers.providers.Provider
-): Promise<T> {
+): Promise<BaseContract> {
   // Generate the contract object that we're going to attach our fancy functions to. Doing it this
   // way is nice because it "feels" more like a contract (as long as you're using ethers).
-  const contract = new ethers.Contract(address, contractInterface, provider) as T;
+  const contract = new ethers.Contract(address, contractInterface, provider);
 
   // Set some code into the contract address so hardhat recognize it as a contract
   await vm.getManager().putContractCode(toFancyAddress(contract.address), Buffer.from('00', 'hex'));
@@ -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..2f95cfa 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -34,31 +34,37 @@ 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;
 }
 
-export type FakeContract<Type extends BaseContract> = BaseContract &
-  Type &
-  {
-    [Property in keyof Type['functions']]: ProgrammableContractFunction;
-  } & {
+export type SmockContractBase<T extends BaseContract> = Omit<BaseContract, 'connect'> &
+  Omit<T, 'connect'> & {
     wallet: Signer;
     fallback: ProgrammableContractFunction;
   };
 
-export type MockContract<Contract extends BaseContract> = BaseContract &
-  Contract &
-  {
-    [Property in keyof Contract['functions']]: ProgrammableContractFunction;
-  } & {
-    wallet: Signer;
-    fallback: ProgrammableContractFunction;
-    setVariable: EditableStorageLogic['setVariable'];
+export type FakeContract<T extends BaseContract = BaseContract> = SmockContractBase<T> & {
+  connect: (...args: Parameters<T['connect']>) => FakeContract<T>;
+} & {
+    [Property in keyof T['functions']]: ProgrammableContractFunction;
+  };
+
+export type MockContract<T extends BaseContract = BaseContract> = SmockContractBase<T> & {
+  connect: (...args: Parameters<T['connect']>) => MockContract<T>;
+  setVariable: EditableStorageLogic['setVariable'];
+} & {
+    [Property in keyof T['functions']]: ProgrammableContractFunction;
   };
 
 type ThenArg<T> = T extends PromiseLike<infer U> ? U : T;


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');
+    });
+  });
+});