Repository: defi-wonderland/smock
Pull request #59 (2021-08-30 10:24:36)

feat: whenCalledWith
**Description**
Revert or return a pre-programmed value based on the call arguments


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