Add ability to mock multiple return values

Ability for the same method to be mocked with a different return value
each time can be useful in mocking utility methods, selector classes,
etc. We use a separate map to keep track of mocks based on method call
count so in the future we can potentially extend this algorithm to
support additional conditions for mocking

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

@@ -12,6 +12,7 @@ public with sharing class UniversalMocker implements System.StubProvider {
   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 {
@@ -177,14 +193,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;
     }
@@ -234,11 +249,25 @@ 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;
+    } else {
+      this.reset();
+    }
+  }
+
+  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) {
@@ -280,7 +309,7 @@ public with sharing class UniversalMocker implements System.StubProvider {
     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 }));
     }
   }
 
@@ -318,6 +347,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>

force-app/main/default/classes/example/AccountDomainTest.cls

@@ -233,6 +233,112 @@ public with sharing class AccountDomainTest {
     System.assertNotEquals(null, acct.Id, 'Account Id is null after insert');
   }
 
+  @IsTest
+  public static void it_should_handle_multiple_return_values_basic() {
+    //setup
+    String mockedMethodName = 'getOneAccount';
+    Account mockAccountOne = new Account(Name = 'Mock Account One');
+    Account mockAccountTwo = new Account(Name = 'Mock Account Two');
+
+    mockService.when(mockedMethodName).thenReturnUntil(1, mockAccountOne).thenReturn(mockAccountTwo);
+
+    //test
+    Test.startTest();
+    Account accountDetail = sut.getAccountDetail();
+    Assert.areEqual(mockAccountOne.Name, accountDetail.Name);
+
+    accountDetail = sut.getAccountDetail();
+    Assert.areEqual(mockAccountTwo.Name, accountDetail.Name);
+
+    //should return mockAccountTwo for all subsequent calls
+    for (Integer i = 0; i < 100; i++) {
+      accountDetail = sut.getAccountDetail();
+    }
+    Assert.areEqual(mockAccountTwo.Name, accountDetail.Name);
+    Test.stopTest();
+
+    //verify
+    mockService.assertThat().method(mockedMethodName).wasCalled(102);
+  }
+
+  @IsTest
+  public static void it_should_handle_multiple_return_values_advanced() {
+    //setup
+    String mockedMethodName = 'getOneAccount';
+    Account mockAccountOne = new Account(Name = 'Mock Account One');
+    Account mockAccountTwo = new Account(Name = 'Mock Account Two');
+    Account mockAccountThree = new Account(Name = 'Mock Account Three');
+
+    mockService.when(mockedMethodName).thenReturnUntil(1, mockAccountOne).thenReturnUntil(3, mockAccountTwo).thenReturn(mockAccountThree);
+
+    //test
+    Test.startTest();
+    Account accountDetail = sut.getAccountDetail();
+    Assert.areEqual(mockAccountOne.Name, accountDetail.Name);
+
+    accountDetail = sut.getAccountDetail();
+    Assert.areEqual(mockAccountTwo.Name, accountDetail.Name);
+
+    accountDetail = sut.getAccountDetail();
+    Assert.areEqual(mockAccountTwo.Name, accountDetail.Name);
+
+    accountDetail = sut.getAccountDetail();
+    Assert.areEqual(mockAccountThree.Name, accountDetail.Name);
+
+    //should return mockAccountTwo for all subsequent calls
+    for (Integer i = 0; i < 100; i++) {
+      accountDetail = sut.getAccountDetail();
+    }
+    Assert.areEqual(mockAccountThree.Name, accountDetail.Name);
+    Test.stopTest();
+
+    //verify
+    mockService.assertThat().method(mockedMethodName).wasCalled(104);
+  }
+
+  @IsTest
+  public static void it_should_handle_multiple_return_values_exception() {
+    //setup
+    String mockedMethodName = 'getOneAccount';
+    Account mockAccountOne = new Account(Name = 'Mock Account One');
+
+    String mockExceptionMessage = 'Mock exception';
+    AuraHandledException mockException = new AuraHandledException(mockExceptionMessage);
+    mockException.setMessage(mockExceptionMessage);
+
+    mockService.when(mockedMethodName).thenThrowUntil(2, mockException).thenReturn(mockAccountOne);
+
+    //test
+    Test.startTest();
+
+    try {
+      Account accountDetail = sut.getAccountDetail();
+      Assert.fail('Expected exception to be thrown');
+    } catch (AuraHandledException ex) {
+      Assert.areEqual(mockExceptionMessage, ex.getMessage());
+    }
+
+    try {
+      Account accountDetail = sut.getAccountDetail();
+      Assert.fail('Expected exception to be thrown');
+    } catch (AuraHandledException ex) {
+      Assert.areEqual(mockExceptionMessage, ex.getMessage());
+    }
+
+    Account accountDetail = sut.getAccountDetail();
+    Assert.areEqual(mockAccountOne.Name, accountDetail.Name);
+
+    //should return mockAccountTwo for all subsequent calls
+    for (Integer i = 0; i < 100; i++) {
+      accountDetail = sut.getAccountDetail();
+    }
+    Assert.areEqual(mockAccountOne.Name, accountDetail.Name);
+    Test.stopTest();
+
+    //verify
+    mockService.assertThat().method(mockedMethodName).wasCalled(103);
+  }
+
   @IsTest
   public static void dummy_test_for_db_service() {
     AccountDBService dbSvc = new AccountDBService();