@@ -56,9 +56,24 @@ public virtual inherited sharing class SOQLCache implements Cacheable {
5656 return null ;
5757 }
5858
59+ public static Cacheable of (SObjectType ofObject ) {
60+ return new SOQLCache (ofObject );
61+ }
62+
63+ public static Cacheable of (String ofObject ) {
64+ return new SOQLCache (ofObject );
65+ }
66+
67+ // Mocking
68+
69+ public interface Mockable {
70+ void thenReturn (SObject record );
71+ }
72+
5973 @TestVisible
60- private static void setMock (String mockId , SObject record ) {
61- mock .setMock (mockId , new List <SObject >{ record });
74+ private static Mockable mock (String mockId ) {
75+ queryIdToMock .put (mockId , new SoqlMock ());
76+ return queryIdToMock .get (mockId );
6277 }
6378
6479 public static void removeFromCache (List <SObject > records ) {
@@ -74,17 +89,16 @@ public virtual inherited sharing class SOQLCache implements Cacheable {
7489 new CacheStorageProxy (ofObject ).sessionCache ().removeRecordsFromCache (records );
7590 }
7691
77- public static Cacheable of (SObjectType ofObject ) {
78- return new SOQLCache (ofObject );
79- }
92+ // Backward support
8093
81- public static Cacheable of (String ofObject ) {
82- return new SOQLCache (ofObject );
94+ @TestVisible // deprecated
95+ private static void setMock (String mockId , SObject record ) {
96+ mock (mockId ).thenReturn (record );
8397 }
8498
8599 // Implementation
86100
87- private static Mock mock = new Mock ();
101+ private static Map < String , SoqlMock > queryIdToMock = new Map < String , SoqlMock > ();
88102
89103 private Executor executor ;
90104 private Cache cache ;
@@ -93,7 +107,7 @@ public virtual inherited sharing class SOQLCache implements Cacheable {
93107 public SOQL.Queryable currentyQuery = null ;
94108
95109 protected SOQLCache (SObjectType ofObject ) {
96- this (ofObject . getDescribe (). getName () );
110+ this (ofObject + ' ' );
97111 }
98112
99113 protected SOQLCache (String ofObject ) {
@@ -126,24 +140,27 @@ public virtual inherited sharing class SOQLCache implements Cacheable {
126140 return this ;
127141 }
128142
129- public Cacheable with (SObjectField field1 ) {
130- return this .with (new List <SObjectField >{ field1 });
143+ public Cacheable with (SObjectField field ) {
144+ this .initialQuery ?. with (field );
145+ this .currentyQuery .with (field );
146+ this .cache .fields .with (field );
147+ return this ;
131148 }
132149
133150 public Cacheable with (SObjectField field1 , SObjectField field2 ) {
134- return this .with (new List < SObjectField >{ field1 , field2 } );
151+ return this .with (field1 ). with ( field2 );
135152 }
136153
137154 public Cacheable with (SObjectField field1 , SObjectField field2 , SObjectField field3 ) {
138- return this .with (new List < SObjectField >{ field1 , field2 , field3 } );
155+ return this .with (field1 , field2 ). with ( field3 );
139156 }
140157
141158 public Cacheable with (SObjectField field1 , SObjectField field2 , SObjectField field3 , SObjectField field4 ) {
142- return this .with (new List < SObjectField >{ field1 , field2 , field3 , field4 } );
159+ return this .with (field1 , field2 , field3 ). with ( field4 );
143160 }
144161
145162 public Cacheable with (SObjectField field1 , SObjectField field2 , SObjectField field3 , SObjectField field4 , SObjectField field5 ) {
146- return this .with (new List < SObjectField >{ field1 , field2 , field3 , field4 , field5 } );
163+ return this .with (field1 , field2 , field3 , field4 ). with ( field5 );
147164 }
148165
149166 public Cacheable with (List <SObjectField > fields ) {
@@ -161,7 +178,7 @@ public virtual inherited sharing class SOQLCache implements Cacheable {
161178 }
162179
163180 public Cacheable whereEqual (SObjectField field , Object value ) {
164- return this .whereEqual (field . getDescribe (). getName () , value );
181+ return this .whereEqual (field + ' ' , value );
165182 }
166183
167184 public Cacheable whereEqual (String field , Object value ) {
@@ -184,7 +201,7 @@ public virtual inherited sharing class SOQLCache implements Cacheable {
184201
185202 public Cacheable mockId (String queryIdentifier ) {
186203 this .currentyQuery .mockId (queryIdentifier );
187- this .executor .mockId ( queryIdentifier );
204+ this .executor .mock ( queryIdToMock . get ( queryIdentifier ) );
188205 return this ;
189206 }
190207
@@ -448,9 +465,13 @@ public virtual inherited sharing class SOQLCache implements Cacheable {
448465 }
449466 }
450467
468+ public void with (SObjectField field ) {
469+ this .cachedFields .add (field + ' ' );
470+ }
471+
451472 public void with (List <SObjectField > fields ) {
452473 for (SObjectField field : fields ) {
453- this .cachedFields .add (field . getDescribe (). getName () );
474+ this .cachedFields .add (field + ' ' );
454475 }
455476 }
456477
@@ -495,27 +516,23 @@ public virtual inherited sharing class SOQLCache implements Cacheable {
495516 }
496517 }
497518
498- private class Mock {
499- private final Map <String , List <SObject >> sObjectsMocks = new Map <String , List <SObject >>();
500-
501- public void setMock (String mockId , List <SObject > records ) {
502- sObjectsMocks .put (mockId , records );
503- }
519+ public class SoqlMock implements Mockable {
520+ private List <SObject > mockedRecords = null ;
504521
505- public Boolean hasMock ( String mockId ) {
506- return sObjectsMocks . containsKey ( mockId ) ;
522+ public void thenReturn ( SObject record ) {
523+ this . mockedRecords = new List < SObject >{ record } ;
507524 }
508525
509- public List <SObject > getSObjectsMock ( String mockId ) {
510- return sObjectsMocks . get ( mockId ) ;
526+ public List <SObject > getMockedResult ( ) {
527+ return mockedRecords ;
511528 }
512529 }
513530
514531 private inherited sharing class Executor {
515532 private SOQL.Queryable currentyQuery ;
516533 private Cache cache ;
517534 private AccessType accessType = null ;
518- private String mockId ;
535+ private SoqlMock mock ;
519536
520537 public Executor (SOQL.Queryable currentyQuery , Cache cache ) {
521538 this .currentyQuery = currentyQuery ;
@@ -526,8 +543,8 @@ public virtual inherited sharing class SOQLCache implements Cacheable {
526543 accessType = type ;
527544 }
528545
529- public void mockId ( String id ) {
530- mockId = id ;
546+ public void mock ( SoqlMock mock ) {
547+ this . mock = mock ;
531548 }
532549
533550 public SObject toObject () {
@@ -545,8 +562,8 @@ public virtual inherited sharing class SOQLCache implements Cacheable {
545562 }
546563
547564 private List <SObject > toList () {
548- if (mock . hasMock ( mockId ) ) {
549- return mock .getSObjectsMock ( mockId );
565+ if (this . mock != null ) {
566+ return this . mock .getMockedResult ( );
550567 }
551568
552569 List <SObject > records = this .cache .toList ();
0 commit comments