Sequential mocks new

README.md

@@ -41,6 +41,35 @@ A universal mocking class for Apex, built using the [Apex Stub API](https://deve
   AccountDBService mockDBService = (AccountDBService)mockInstance.createStub();
   ```
 
+#### Sequential Mocks
+
+There might be instances where you may need the same method to mock different return values within the same test when
+testing utility methods or selector classes and such. You can specify different return values based on the call count
+in such cases
+
+- Basic example
+
+```java
+  mockInstance.when('getOneAccount').thenReturnUntil(3,mockAccountOne).thenReturn(mockAccountTwo);
+```
+
+Here, `mockAccountOne` is returned the first 3 times `getOneAccount` is called. All subsequent calls to `getOneAccount`
+will return `mockAccountTwo`
+
+- You can also pair it with param types or to mock exceptions
+
+```java
+  mockInstance.when('getOneAccount').withParamTypes(new List<Type>{Id.class})
+    .thenReturnUntil(1,mockAccountOne)
+    .thenThrowUntil(3,mockException)
+    .thenReturn(mockAccountTwo);
+```
+
+Refer to the [relevant unit tests](force-app/main/default/classes/example/AccountDomainTest.cls#L265) for further
+clarity
+
+**Note**: It is recommended that you end all setup method call chains with `thenReturn` or `thenThrow`
+
 #### Mutating arguments
 
 There might be instances where you need to modify the original arguments passed into the function. A typical example

force-app/main/default/classes/UniversalMocker.cls

@@ -6,12 +6,13 @@
 *** @description: A universal class for mocking in tests. Contains a method for setting the return value for any method. Another method returns the number of times a method was called. https://github.com/surajp/universalmock
 
 */
-@isTest
+@IsTest
 public with sharing class UniversalMocker implements System.StubProvider {
   // Map of methodName+paramTypes -> map of (paramname,value) for each invocation
   private final Map<String, List<Map<String, Object>>> argumentsMap = new Map<String, List<Map<String, Object>>>();
   private final Type mockedClass;
   private final Map<String, Object> mocksMap = new Map<String, Object>();
+  private final Map<String, List<Integer>> returnUntilMap = new Map<String, List<Integer>>();
   private final Map<String, Integer> callCountsMap = new Map<String, Integer>();
 
   @TestVisible
@@ -24,6 +25,7 @@ public with sharing class UniversalMocker implements System.StubProvider {
   private String currentParamTypesString;
   private Integer expectedCallCount;
   private Integer forInvocationNumber = 0;
+  private Integer callCountToMock = null;
 
   private String KEY_DELIMITER = '||';
 
@@ -70,9 +72,19 @@ public with sharing class UniversalMocker implements System.StubProvider {
     }
   }
 
-  public virtual class IntermediateSetupState {
-    private final UniversalMocker parent;
+  public virtual class IntermediateSetupState extends FinalSetupState {
     private IntermediateSetupState(UniversalMocker parent) {
+      super(parent);
+    }
+    public FinalSetupState mutateWith(Mutator mutatorInstance) {
+      this.parent.mutateWith(mutatorInstance);
+      return (FinalSetupState) this;
+    }
+  }
+
+  public virtual class FinalSetupState {
+    private final UniversalMocker parent;
+    private FinalSetupState(UniversalMocker parent) {
       this.parent = parent;
     }
     public void thenReturnVoid() {
@@ -81,13 +93,17 @@ public with sharing class UniversalMocker implements System.StubProvider {
     public void thenReturn(Object returnObject) {
       this.parent.thenReturn(returnObject);
     }
-    public IntermediateSetupState mutateWith(Mutator mutatorInstance) {
-      this.parent.mutateWith(mutatorInstance);
-      return this;
-    }
     public void thenThrow(Exception exceptionToThrow) {
       this.parent.thenThrow(exceptionToThrow);
     }
+    public FinalSetupState thenReturnUntil(Integer callCount, Object returnObject) {
+      this.parent.thenReturnUntil(callCount, returnObject);
+      return this;
+    }
+    public FinalSetupState thenThrowUntil(Integer callCount, Exception exceptionToThrow) {
+      this.parent.thenThrowUntil(callCount, exceptionToThrow);
+      return this;
+    }
   }
 
   public class InitialValidationState {
@@ -161,6 +177,7 @@ public with sharing class UniversalMocker implements System.StubProvider {
   }
 
   public InitialSetupState when(String stubbedMethodName) {
+    this.reset();
     this.currentMethodName = stubbedMethodName;
     return this.setupAInstance;
   }
@@ -177,14 +194,13 @@ public with sharing class UniversalMocker implements System.StubProvider {
     this.incrementCallCount(keyInUse);
     this.saveArguments(listOfParamNames, listOfArgs, keyInUse);
 
-    Object returnValue = this.mocksMap.get(keyInUse);
-
     if (this.mutatorMap.containsKey(keyInUse)) {
       for (Mutator m : this.mutatorMap.get(keyInUse)) {
         m.mutate(stubbedObject, stubbedMethodName, listOfParamTypes, listOfArgs);
       }
     }
 
+    Object returnValue = this.getMockValue(keyInUse);
     if (returnValue instanceof Exception) {
       throw (Exception) returnValue;
     }
@@ -193,10 +209,12 @@ public with sharing class UniversalMocker implements System.StubProvider {
   }
 
   public InitialValidationState assertThat() {
+    this.reset();
     return this.assertAInstance;
   }
 
   public InitialParamValidationState forMethod(String stubbedMethodName) {
+    this.reset();
     this.currentMethodName = stubbedMethodName;
     return this.getParamsAInstance;
   }
@@ -234,11 +252,23 @@ public with sharing class UniversalMocker implements System.StubProvider {
 
   private void thenReturn(Object returnObject) {
     String key = this.getCurrentKey();
-    this.mocksMap.put(key, returnObject);
+    this.putMockValue(key, returnObject);
     if (!this.callCountsMap.containsKey(key)) {
       this.callCountsMap.put(key, 0);
     }
-    this.reset();
+    if (this.callCountToMock != null) {
+      this.callCountToMock = null;
+    }
+  }
+
+  private void thenReturnUntil(Integer callCount, Object returnObject) {
+    this.callCountToMock = callCount;
+    this.thenReturn(returnObject);
+  }
+
+  private void thenThrowUntil(Integer callCount, Exception exceptionToThrow) {
+    this.callCountToMock = callCount;
+    this.thenReturn(exceptionToThrow);
   }
 
   private void thenThrow(Exception exceptionToThrow) {
@@ -259,7 +289,6 @@ public with sharing class UniversalMocker implements System.StubProvider {
     //Integer actualCallCount = this.callCountsMap.get(currentKey);
     Integer actualCallCount = this.getCallCountsMapInternal().get(currentKey);
     String methodName = this.currentMethodName;
-    this.reset();
     switch on assertTypeValue {
       when OR_LESS {
         system.assert(this.expectedCallCount >= actualCallCount, this.getMethodCallCountAssertMessage(methodName, 'less than or equal'));
@@ -277,10 +306,9 @@ public with sharing class UniversalMocker implements System.StubProvider {
     String currentKey = this.getCurrentKey();
     Integer actualCallCount = this.getCallCountsMapInternal().get(currentKey);
     String methodName = this.currentMethodName;
-    this.reset();
     if (actualCallCount != null) {
       this.expectedCallCount = 0;
-      system.assertEquals(this.expectedCallCount, actualCallCount, String.format('Method {0} was called 1 or more times', new List<String>{ methodName }));
+      System.assertEquals(this.expectedCallCount, actualCallCount, String.format('Method {0} was called 1 or more times', new List<String>{ methodName }));
     }
   }
 
@@ -295,14 +323,12 @@ public with sharing class UniversalMocker implements System.StubProvider {
       throw new IllegalArgumentException(String.format('Param name {0} not found for the method {1}', new List<Object>{ paramName, this.currentMethodName }));
     }
     Object returnValue = paramsMap.get(paramName.toLowerCase());
-    this.reset();
     return returnValue;
   }
 
   private Map<String, Object> getArgumentsMap() {
     String theKey = this.getCurrentKey();
     Map<String, Object> returnValue = this.getArgumentsMapInternal().get(theKey).get(this.forInvocationNumber);
-    this.reset();
     return returnValue;
   }
 
@@ -318,6 +344,32 @@ public with sharing class UniversalMocker implements System.StubProvider {
     return (methodName + KEY_DELIMITER + this.getParamTypesString(paramTypes)).toLowerCase();
   }
 
+  private Object getMockValue(String key) {
+    if (this.returnUntilMap.containsKey(key)) {
+      Integer callCount = this.callCountsMap.get(key);
+      List<Integer> returnUntilList = this.returnUntilMap.get(key);
+      returnUntilList.sort();
+      for (Integer returnUntil : returnUntilList) {
+        if (returnUntil >= callCount) {
+          return this.mocksMap.get(key + '-' + returnUntil);
+        }
+      }
+    }
+    return this.mocksMap.get(key);
+  }
+
+  private void putMockValue(String key, Object value) {
+    if (this.callCountToMock != null) {
+      if (!this.returnUntilMap.containsKey(key)) {
+        this.returnUntilMap.put(key, new List<Integer>{});
+      }
+      this.returnUntilMap.get(key).add(this.callCountToMock);
+      this.mocksMap.put(key + '-' + this.callCountToMock, value);
+    } else {
+      this.mocksMap.put(key, value);
+    }
+  }
+
   private String getParamTypesString(List<Type> paramTypes) {
     String[] classNames = new List<String>{};
     for (Type paramType : paramTypes) {

force-app/main/default/classes/UniversalMocker.cls-meta.xml

@@ -1,5 +1,5 @@
 <?xml version="1.0" encoding="UTF-8" ?>
 <ApexClass xmlns="http://soap.sforce.com/2006/04/metadata">
-    <apiVersion>56.0</apiVersion>
+    <apiVersion>59.0</apiVersion>
     <status>Active</status>
 </ApexClass>