Make mutator behavior consistent with return value behavior

`Mutatewith` mutators will only be applied after all previous
`mutateUntil` calls are done. Mutators with the same value of
`mutateUntil` will be accumulated and applied in succession. Also
updated README.

README.md

@@ -10,6 +10,8 @@ A universal mocking class for Apex, built using the [Apex Stub API](https://deve
 
 ### Setup
 
+#### The Basics
+
 - Create an instance of `UniversalMocker` for each class you want to mock.
 
   ```java
@@ -70,7 +72,7 @@ clarity
 
 **Note**: It is recommended that you end all setup method call chains with `thenReturn` or `thenThrow`
 
-#### Mutating arguments
+#### Mutating Arguments
 
 There might be instances where you need to modify the original arguments passed into the function. A typical example
 would be to set the `Id` field of records passed into a method responsible for inserting them.
@@ -97,19 +99,69 @@ Here's the method for setting fake ids on inserted records, in our example.
   }
 ```
 
-Check out the [AccountDomainTest](./force-app/main/default/classes/example/AccountDomainTest.cls#L187) class for the
-full example.
-
-- Pass in an instance of your implementation of the `Mutator` class to mutate the method arguments. Check out the
-  complete test method [here](./force-app/main/default/classes/example/AccountDomainTest.cls#L146)
+- Pass in an instance of your implementation of the `Mutator` class to mutate the method arguments.
 
 ```java
   mockInstance.when('doInsert').mutateWith(dmlMutatorInstance).thenReturnVoid();
 ```
 
-**Note**: You can call the `mutateWith` method any number of times in succession, with the same or different mutator instances,
+Check out the [AccountDomainTest](./force-app/main/default/classes/example/AccountDomainTest.cls#L244) class for the
+full example.
+
+You can call the `mutateWith` method any number of times in succession, with the same or different mutator instances,
 to create a chain of methods to mutate method arguments.
 
+#### Sequential Mutators
+
+You can also use specific mutators based on call count. Multiple mutators with the same value of call count will be
+accumulated and applied in succession for all calls since the previous established call count.
+
+For example, lets say you have a `DescriptionMutator` class as shown below. It appends a given string to the `Account
+Description` field.
+
+```java
+  //Adds a given suffix to account description
+  public class DescriptionMutator implements UniversalMocker.Mutator {
+    private String stringToAdd = '';
+    public DescriptionMutator(String stringToAdd) {
+      this.stringToAdd = stringToAdd;
+    }
+    public void mutate(Object stubbedObject, String stubbedMethodName, List<Type> listOfParamTypes, List<Object> listOfArgs) {
+      Account record = (Account) listOfArgs[0];
+      if (record.get('Description') != null) {
+        record.Description += this.stringToAdd;
+      } else {
+        record.Description = this.stringToAdd;
+      }
+    }
+  }
+```
+
+If you wanted to append the string `12` to the Account Description for the first 2 calls and then the string `3` for all subsequent
+calls, your setup would look something like:
+
+```java
+mockService.when(mockedMethodName).mutateUntil(2, new DescriptionMutator('1')).mutateUntil(2, new DescriptionMutator('2'))
+.mutateWith(new DescriptionMutator('3'));
+```
+
+or
+
+```java
+mockService.when(mockedMethodName).mutateUntil(2, new DescriptionMutator('12')).mutateWith(new DescriptionMutator('3'));
+```
+
+If you wanted to append the string `1` to the Account Description for the first call, the string `2` for the second call,
+and string `3` for all subsequent calls, your setup would look as follows:
+
+```java
+mockService.when(mockedMethodName).mutateUntil(1, new DescriptionMutator('1')).mutateUntil(2, new DescriptionMutator('2'))
+.mutateWith(new DescriptionMutator('3'));
+```
+
+Check out the [AccountDomainTest](./force-app/main/default/classes/example/AccountDomainTest.cls#L193) class for the
+full example.
+
 ### Verification
 
 - Assert the exact number of times a method was called.

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

@@ -9,12 +9,12 @@
 @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 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, List<Integer>> mutateUntilMap = new Map<String, List<Integer>>();
-  private final Map<String, Integer> callCountsMap = new Map<String, Integer>();
+  private Map<String, Object> mocksMap = new Map<String, Object>();
+  private Map<String, List<Integer>> returnUntilMap = new Map<String, List<Integer>>();
+  private Map<String, List<Integer>> mutateUntilMap = new Map<String, List<Integer>>();
+  private Map<String, Integer> callCountsMap = new Map<String, Integer>();
 
   @TestVisible
   private static final Map<String, UniversalMocker> uMockInstances = new Map<String, UniversalMocker>();
@@ -24,7 +24,6 @@ public with sharing class UniversalMocker implements System.StubProvider {
 
   private String currentMethodName;
   private String currentParamTypesString;
-  private Integer expectedCallCount;
   private Integer forInvocationNumber = 0;
   private Integer callCountToMock = null;
 
@@ -223,6 +222,17 @@ public with sharing class UniversalMocker implements System.StubProvider {
     void mutate(Object stubbedObject, String stubbedMethodName, List<Type> listOfParamTypes, List<Object> listOfArgs);
   }
 
+  public void resetState() {
+    this.reset();
+    this.argumentsMap = new Map<String, List<Map<String, Object>>>();
+    this.mocksMap = new Map<String, Object>();
+    this.returnUntilMap = new Map<String, List<Integer>>();
+    this.mutateUntilMap = new Map<String, List<Integer>>();
+    this.callCountsMap = new Map<String, Integer>();
+    this.mutatorMap = new Map<String, List<Mutator>>();
+    this.initInnerClassInstances();
+  }
+
   /* End Public methods */
 
   /* Begin Private methods */
@@ -237,6 +247,9 @@ public with sharing class UniversalMocker implements System.StubProvider {
     if (!this.callCountsMap.containsKey(key)) {
       this.callCountsMap.put(key, 0);
     }
+    if (this.callCountToMock != null) {
+      this.callCountToMock = null;
+    }
   }
 
   private void thenReturnVoid() {
@@ -282,20 +295,19 @@ public with sharing class UniversalMocker implements System.StubProvider {
   }
 
   private void wasCalled(Integer expectedCallCount, Times assertTypeValue) {
-    this.expectedCallCount = expectedCallCount;
     String currentKey = this.getCurrentKey();
     //Integer actualCallCount = this.callCountsMap.get(currentKey);
     Integer actualCallCount = this.getCallCountsMapInternal().get(currentKey);
     String methodName = this.currentMethodName;
     switch on assertTypeValue {
       when OR_LESS {
-        system.assert(this.expectedCallCount >= actualCallCount, this.getMethodCallCountAssertMessage(methodName, 'less than or equal'));
+        system.assert(expectedCallCount >= actualCallCount, this.getMethodCallCountAssertMessage(methodName, 'less than or equal'));
       }
       when OR_MORE {
-        system.assert(this.expectedCallCount <= actualCallCount, this.getMethodCallCountAssertMessage(methodName, 'more than or equal'));
+        system.assert(expectedCallCount <= actualCallCount, this.getMethodCallCountAssertMessage(methodName, 'more than or equal'));
       }
       when else {
-        system.assertEquals(this.expectedCallCount, actualCallCount, this.getMethodCallCountAssertMessage(methodName, 'equal'));
+        system.assertEquals(expectedCallCount, actualCallCount, this.getMethodCallCountAssertMessage(methodName, 'equal'));
       }
     }
   }
@@ -305,8 +317,8 @@ public with sharing class UniversalMocker implements System.StubProvider {
     Integer actualCallCount = this.getCallCountsMapInternal().get(currentKey);
     String methodName = this.currentMethodName;
     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 }));
+      Integer expectedCallCount = 0;
+      System.assertEquals(expectedCallCount, actualCallCount, String.format('Method {0} was called 1 or more times', new List<String>{ methodName }));
     }
   }
 

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

@@ -14,4 +14,8 @@ public with sharing class AccountDBService {
   public void doInsert(Account acct) {
     insert acct;
   }
+
+  public void doUpdate(Account acct) {
+    update acct;
+  }
 }

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

@@ -18,6 +18,10 @@ public with sharing class AccountDomain {
     this.acctService.doInsert(acct);
   }
 
+  public void updateAccount(Account acct) {
+    this.acctService.doUpdate(acct);
+  }
+
   public Account[] getMatchingAccounts(String attribute) {
     if (attribute instanceof Id) {
       return this.acctService.getMatchingAccounts(Id.valueOf(attribute));

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

@@ -189,6 +189,40 @@ public with sharing class AccountDomainTest {
     System.assertNotEquals(null, acct.Id, 'Account Id is null after insert');
   }
 
+  @IsTest
+  public static void shouldApplyMutatorsBasedOnCallCounts() {
+    String mockedMethodName = 'doUpdate';
+    String mockExceptionMessage = 'Mock exception';
+    UniversalMocker.Mutator dmlMutatorInstance = new DMLMutator();
+
+    mockService.when(mockedMethodName).mutateUntil(1, new DescriptionMutator('1')).thenReturnVoid();
+
+    Account acct = new Account(Name = 'Acme Inc');
+    sut.updateAccount(acct);
+    sut.updateAccount(acct);
+    sut.updateAccount(acct);
+
+    //verify
+    mockService.assertThat().method(mockedMethodName).wasCalled(3);
+    Assert.areEqual('1', acct.Description, 'Expected description to only be set once (till callcount 1)');
+
+    mockService.resetState();
+    acct.Description = ''; //reset account description
+    mockService.when(mockedMethodName)
+      .mutateUntil(2, new DescriptionMutator('1'))
+      .mutateUntil(2, new DescriptionMutator('2'))
+      .mutateWith(new DescriptionMutator('3'))
+      .thenReturnVoid();
+    sut.updateAccount(acct);
+    sut.updateAccount(acct);
+    sut.updateAccount(acct);
+    Assert.areEqual(
+      '12123',
+      acct.Description,
+      'Expected description to set to "12123" ("12" appended for each of the first two calls and "3" appended for the third call'
+    );
+  }
+
   @IsTest
   public static void it_should_track_call_counts_with_batchables() {
     String mockedMethodName = 'getOneAccount';
@@ -399,6 +433,22 @@ public with sharing class AccountDomainTest {
     dbSvc.getMatchingAccounts('Acme');
   }
 
+  //Adds a given suffix to account description
+  public class DescriptionMutator implements UniversalMocker.Mutator {
+    private String stringToAdd = '';
+    public DescriptionMutator(String stringToAdd) {
+      this.stringToAdd = stringToAdd;
+    }
+    public void mutate(Object stubbedObject, String stubbedMethodName, List<Type> listOfParamTypes, List<Object> listOfArgs) {
+      Account record = (Account) listOfArgs[0];
+      if (record.get('Description') != null) {
+        record.Description += this.stringToAdd;
+      } else {
+        record.Description = this.stringToAdd;
+      }
+    }
+  }
+
   public class DMLMutator implements UniversalMocker.Mutator {
     // Ideally, 'fakeCounter' should be a static variable and 'getFakeId' should be a static method in another top-level class.
     private Integer fakeIdCounter = 1;