Mocking Database.QueryLocator in Unit Tests using fflib Apex Mocks

Question

Assume we have some basic AccountSelector class that has different implementations of the selector method. QueryLocator variation is helpful once we need to take advantage of SOQL for loops or to use it in the start() method of batches.

public with sharing class AccountSelector extends fflib_SObjectSelector {

    public override SObjectType getSObjectType() {
        return Account.SObjectType;
    }

    public override List<SObjectField> getSObjectFieldList() {
        return new List<SObjectField>{ 
            Account.Id,
            Account.Name
        };
    }

    public List<Account> byIds(Set<Id> ids) {
        return (List<Account>) Database.query(soqlByIds(ids));
    }

    public Database.QueryLocator locatorByIds(Set<Id> ids) {
        return Database.getQueryLocator(soqlByIds(ids));
    }

    private String soqlByIds(Set<Id> ids) {
        fflib_QueryFactory queryFactory = newQueryFactory()
                .setEnforceFLS(true)
                .setCondition('Id IN :ids');
        return queryFactory.toSOQL();
    }
}

Some service class uses AccountSelector and I need to unit test that service class using fflib Apex Mocks:

@IsTest(IsParallel=true)
private class SomeServiceUnitTest {

    @IsTest
    static void shouldReturnAccount() {
        // Create Mocks
        fflib_ApexMocks mocks = new fflib_ApexMocks();
        AccountSelector accSelectorMock = (AccountSelector) mocks.mock(AccountSelector.class);
        // Given
        List<Account> accountsMock = new List<Account>();
        Id accId = fflib_IDGenerator.generate(Account.SObjectType);
        accountsMock.add(
                new Account(
                        Id = accId,
                        Name = 'TestName'
                )
        );
        mocks.startStubbing();
            mocks.when(accSelectorMock.sObjectType())
                .thenReturn(Account.SObjectType);
            mocks.when(accSelectorMock.byIds(new Set<Id>{ accId }))
                .thenReturn(accountsMock); // Works as expected
            mocks.when(accSelectorMock.locatorByIds(new Set<Id>{ accId }))
                .thenReturn(???); // Need to get/create an instance of Database.QueryLocator to mock
        mocks.stopStubbing();
        Application.SELECTOR.setMock(accSelectorMock);
    }

}

So the problem here is that when we want to mock accSelectorMock.locatorByIds() method invocation, we need to get an instance of Database.QueryLocator. The known way to get it is Database.getQueryLocator() static method which can take either String query or List<SObject> query param e.g:

Database.getQueryLocator('SELECT Id FROM Account'); 
Database.getQueryLocator([SELECT Id FROM Account]);

Database.QueryLocator instance represents a server-side database cursor but in case if we want to write pure unit tests without any DML operations(e.g. inserting test accounts) it seems impossible to get a mock of Database.QueryLocator instance somehow. So the questions are:

Can we get an instance of Database.QueryLocator that points to some test records without first inserting them?

How to mock a selector method that returns Database.QueryLocator in the unit tests in terms of fflib Apex Mocks?

Answer 1

Currently, I came up with some ugly but efficient solution for mocking Database.QueryLocator and to be precise, its substitution.

1. To change the return type of AccountSelector.locatorByIds() method from Database.QueryLocator to Object:

public with sharing class AccountSelector extends fflib_SObjectSelector {
    public override SObjectType getSObjectType() {
        return Account.SObjectType;
    }

    public override List<SObjectField> getSObjectFieldList() {
        return new List<SObjectField>{ 
            Account.Id,
            Account.Name
        };
    }

    public List<Account> byIds(Set<Id> ids) {
        return (List<Account>) Database.query(soqlByIds(ids)));
    }

    public Object locatorByIds(Set<Id> ids) {
        return Database.getQueryLocator(soqlByIds(ids));
    }

    private String soqlByIds(Set<Id> ids) {
        fflib_QueryFactory queryFactory = newQueryFactory()
                .setEnforceFLS(true)
                .setCondition('Id IN :ids');
        return queryFactory.toSOQL();
    }
    
}

I also tried

    public Iterable<Account> locatorByIds(Set<Id> ids) {
        return (Iterable<Account>) Database.getQueryLocator(soqlByIds(ids));
    }

and

    public Iterator<Account> locatorByIds(Set<Id> ids) {
        return (Iterator<Account>) Database.getQueryLocator(soqlByIds(ids)).iterator();
    }

method variations and they work as expected but they have one significant drawback in terms of fflib Apex Mocks - they cannot be mocked due to Stub API Limitations. Every time invoking the mock() method on the fflib_ApexMocks instance with Selector type parameter of the Selector class which has at least 1 method that returns Iterator<T> or Iterable<T> fails with:

fflib_ApexMocks mocks = new fflib_ApexMocks();
AccountSelector accSelectorMock = (AccountSelector) mocks.mock(AccountSelector.class);

System.TypeException: No type arguments provided for a parameterized type: System.Iterable (or System.Iterator) Class.System.Test.createStub: line 93, column 1 Class.fflib_ApexMocks.mock: line 67, column 1

However, there are no limitations to mock the method that returns Object.

2. So now the locatorByIds() method can be mocked next way:

@IsTest(IsParallel=true)
private class SomeServiceUnitTest {

    @IsTest
    static void shouldDoSomething() {
        // Create Mocks
        fflib_ApexMocks mocks = new fflib_ApexMocks();
        AccountSelector accSelectorMock = (AccountSelector) mocks.mock(AccountSelector.class);
        // Given
        List<Account> accountsMock = new List<Account>();
        Id accId = fflib_IDGenerator.generate(Account.SObjectType);
        accountsMock.add(
                new Account(
                        Id = accId,
                        Name = 'TestName'
                )
        );
        mocks.startStubbing();
            mocks.when(accSelectorMock.sObjectType())
                .thenReturn(Account.SObjectType);
            mocks.when(accSelectorMock.locatorByIds(new Set<Id>{ accId }))
                .thenReturn(accountsMock); 
        mocks.stopStubbing();
        Application.SELECTOR.setMock(accSelectorMock);
    }

}

3. To add some dynamic casting directly in the Service (or in another suitable layer):

public with sharing class SomeService {

    public void someMethod(Set<Id> accIds) {
        AccountSelector accSelector =
            (AccountSelector) Application.SELECTOR.newInstance(Account.SObjectType);

        Object accLocatorLike = accSelector.locatorByIds(accIds);
        Iterator<Account> accIterator;
        if (accLocatorLike instanceof Database.QueryLocator) {
            accIterator = (Iterator<Account>) ((Database.QueryLocator) accLocatorLike).iterator();
        } else if (accLocatorLike instanceof List<Account>) {
            accIterator = (Iterator<Account>) ((Iterable<Account>) accLocatorLike).iterator();
        }
        while (accIterator != null && accIterator.hasNext()) {
            Account acc = accIterator.next();
            // Processing the record returned by iterator
        }
        /* OR
        Iterator<SObject> sObjIterator;
        if (accLocatorLike instanceof Database.QueryLocator) {
            sObjIterator = ((Database.QueryLocator) accLocatorLike).iterator();
        } else if (accLocatorLike instanceof List<Account>) {
            sObjIterator = ((Iterable<SObject>) accLocatorLike).iterator();
        }
        while (sObjIterator != null && sObjIterator.hasNext()) {
            Account acc = (Account) sObjIterator.next();
            // Processing the record returned by iterator
        }
        */
    }
}

The main point here is that we should perform dynamic casting for the entity returned by the locatorByIds() method depending on the context. When the returned entity is of type Database.QueryLocator it can be cast to Database.QueryLocator apriori. When the returned entity is of type List<Account> it can be cast to Iterable<Account> because List<T> implements Iterable<T>. So now we have either Database.QueryLocator or Iterable<Account> instance and both have the iterator() method since Database.QueryLocator implements Iterable. Since Database.QueryLocatorIterator implements Iterator both iterators can explicitly be cast to Iterator<Account>.

Although this solution looks a bit ugly, the performance cost of casting is cheap enough in JAVA and AFAIK Apex uses JVM.

PS Hopefully, someone will find a more sophisticated solution. Please write in the comments so that I can improve the answer.

PSS Please vote for Add Generics to the Apex Language everyone who has not yet voted.

Answer 2

I don't think you can mock it. I assume you need this b/c you want to mock out some records to test Batchable classes? If so, I just call each method (start, execute, finish) independently. The records I want to process are passed in manually within the execute method. The first method here will require you to have real records inserted for your tests. The 2nd method you can mock out the records you want to process:

public inherited sharing class SynchronousBatchRunner {
    public static void run(Database.Batchable<SObject> batchableClass) {
        Database.QueryLocator locator = (Database.QueryLocator) batchableClass.start(null);
        Database.QueryLocatorIterator it = locator.iterator();
        List<SObject> records = new List<SObject>();
        while (it.hasNext()) {
            records.add(it.next());
        }
        batchableClass.execute(null, records);
        batchableClass.finish(null);
    }

    public static void run(Database.Batchable<SObject> batchableClass, List<SObject> recordList) {
        batchableClass.start(null);
        batchableClass.execute(null, recordList);
        batchableClass.finish(null);

    }
}

Answer 3

Apologies for being late to this party, but I wanted to share the technique that I have used before. It doesn't cover all use cases, but it does cover some of them.

What I have done in the past is, prior to point where you start to stub out the mock selector, and prior to where you would let the Application Selector Factory know about the mock selector, you could make a call to the real version of the selector (just instantiate the selector directly) method and get a handle to the returning QueryLocator. Then, during the Apex Mocks stubbing phase, you can stub the selector method like normal and setup the QueryLocator object that you got from the real selector method call as the return value from the method call.

Doing that allows you to use Apex Mock's various "verify" methods and keep track of the selector method execution for testing/verification purposes.

Again, this approach doesn't cover all possible testing needs around mocking a selector method with QueryLocator return type, but it does address some.

Hope this helps.

Cheers!