From b45ba98dbc5310ae83da780c8046164fc88116ab Mon Sep 17 00:00:00 2001
From: James Simone <16430727+jamessimone@users.noreply.github.com>
Date: Wed, 24 Apr 2024 14:45:54 -0400
Subject: [PATCH] Adds count() override for mocking and getCount() method on
 AggregateRecord

---
 force-app/factory/RepoFactoryMock.cls        | 14 ++++++++++++++
 force-app/repository/AggregateRecord.cls     | 10 ++++++++++
 force-app/repository/AggregateRepository.cls |  2 +-
 3 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/force-app/factory/RepoFactoryMock.cls b/force-app/factory/RepoFactoryMock.cls
index f1c1eff..eacdbc6 100644
--- a/force-app/factory/RepoFactoryMock.cls
+++ b/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);
diff --git a/force-app/repository/AggregateRecord.cls b/force-app/repository/AggregateRecord.cls
index 1d12051..91049ed 100644
--- a/force-app/repository/AggregateRecord.cls
+++ b/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();
diff --git a/force-app/repository/AggregateRepository.cls b/force-app/repository/AggregateRepository.cls
index d3670f5..8359639 100644
--- a/force-app/repository/AggregateRepository.cls
+++ b/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();