Adds count() override for mocking and getCount() method on AggregateR…

…ecord

force-app/factory/RepoFactoryMock.cls

@@ -126,14 +126,28 @@ public class RepoFactoryMock {
 
     public override List<SObject> get(List<Query> queries) {
       QueriesMade.addAll(queries);
+      this.clearState();
       return this.results;
     }
 
     public override List<List<SObject>> getSosl(String searchTerm, List<Query> queries) {
       QueriesMade.addAll(queries);
+      this.clearState();
       return new List<List<SObject>>{ this.results };
     }
 
+    public override Integer count(List<Query> queries) {
+      QueriesMade.addAll(queries);
+      List<AggregateRecord> results = AggregateResults.get(this.repoType);
+      this.clearState();
+      if (results == null) {
+        return super.count(queries);
+      } else if (results.isEmpty() == false) {
+        return results.remove(0).getCount();
+      }
+      return null;
+    }
+
     public override List<AggregateRecord> aggregate(List<Aggregation> aggregations, List<Query> queries) {
       AggregatesMade.addAll(aggregations);
       QueriesMade.addAll(queries);

force-app/repository/AggregateRecord.cls

@@ -1,5 +1,6 @@
 public class AggregateRecord {
   private final Map<String, Object> keyToAggregateResult = new Map<String, Object>();
+  private static final String COUNT_KEY = 'countKey';
 
   public AggregateRecord putAll(Map<String, Object> values) {
     this.keyToAggregateResult.putAll(values);
@@ -10,6 +11,15 @@ public class AggregateRecord {
     return this.keyToAggregateResult.get(key);
   }
 
+  public Integer getCount() {
+    return (Integer) this.keyToAggregateResult.get(COUNT_KEY);
+  }
+
+  public AggregateRecord setCount(Integer countAmount) {
+    this.keyToAggregateResult.put(COUNT_KEY, countAmount);
+    return this;
+  }
+
   public Boolean equals(Object that) {
     if (that instanceof AggregateResult) {
       Map<String, Object> thatKeyToAggregateResult = ((AggregateResult) that).getPopulatedFieldsAsMap();

force-app/repository/AggregateRepository.cls

@@ -91,7 +91,7 @@ public without sharing virtual class AggregateRepository extends Repository impl
   public Integer count(Query query) {
     return this.count(new List<Query>{ query });
   }
-  public Integer count(List<Query> queries) {
+  public virtual Integer count(List<Query> queries) {
     this.isNumberCountQuery = true;
     Integer recordCount = Database.countQueryWithBinds(this.getFinalQuery(queries), this.bindVars, this.accessLevel);
     this.clearState();