Flexible Mock

Flexible Mock

TL; DR;

Simplifying Apex Mocking with a flexible Stub class which can replace any “mockable” method call and argument combination.

See on GitHub

The Problem

Unit tests are hard. At least good unit tests are. So often you see repeating complicated setup code which creates dozens of dependent records so that the code baked into triggers can actually be even executed. And because of that all the tests are so slow you only ever run them when you really really have to.

One of the ways of improving the above situation is to use Mocking. This allows to not work with the database so much and therefore skip all the time consuming Apex Trigger trips and complicated logic to setup all the required data*. What I want to talk about here is a little test utility class I created after having repeated a similar thing in a few test scenarios.

Being able to do this is not always so simple and requires code to be organised a certain way. We not going to get into that in this post.

Repeating Story

A few times I found myself creating an inner class inside a Test class, which would be initialised by a bunch of tests with lists of records to return when called. I’d then use that to “mock out” a singleton Service class.

Here’s an example:

public class CampaignEmailReminderServiceMock implements System.StubProvider {
    private List<ADvendio__MediaCampaign__c> mockCampaigns;
    private Boolean throwExceptionDuringSend;

    public CampaignEmailReminderServiceMock(List<ADvendio__MediaCampaign__c> mockCampaigns, Boolean throwExceptionDuringSend) {
        this.mockCampaigns = mockCampaigns;
        this.throwExceptionDuringSend = throwExceptionDuringSend;
    }

    public Object handleMethodCall(
        Object stubbedObject,
        String stubbedMethodName,
        Type returnType,
        List<Type> listOfParamTypes,
        List<String> listOfParamNames,
        List<Object> listOfArgs
    ) {
        if ('getMatchingCampaigns'.equals(stubbedMethodName)) {
            return this.mockCampaigns;
        } else if ('sendNotificationEmails'.equals(stubbedMethodName)) {
            if (this.throwExceptionDuringSend) {
                throw new CampaignEmailReminderServiceMockException('Mock Exception');
            } else {
                mockSendEmailsInvocationCount++;
            }
        }
        return null;
    }
}

The real service class starts with a test accessible private static member to hold the mock version. It can be set by Unit Tests that need it and live code calling the public newInstance() method is none the wiser.

public with sharing class CampaignEmailReminderService {

   @TestVisible
   private static CampaignEmailReminderService mockInstance;
   public static CampaignEmailReminderService newInstance(CampaignEmailReminderSetting setting) {
       if (mockInstance != null) {
           return mockInstance;
       }
       return new CampaignEmailReminderService(setting);
   }
...

Each tests can prepare a list of records and make sure that those are returned by the Reminder Service:

CampaignEmailReminderService.mockInstance = (CampaignEmailReminderService) Test.createStub(
    CampaignEmailReminderService.class,
    new CampaignEmailReminderServiceMock(new List<ADvendio__MediaCampaign__c>{ new ADvendio__MediaCampaign__c() }, false)
);

The Solution

As I mentioned, I’ve done something similar a few times and naturally wanted to make it a bit less repetitive. Oftentimes there are many different services or selectors contributing input to the tested part of the application. They all need to be mocked to get the full benefit from the approach.

So I built myself a reusable class for mocking different methods. It allows me to register multiple methodName and argument combinations and what I would like returned given those combinations.

The return type is Object so as to be able to mock anything. I’m hashing the argument lists and using that as the key in a map of results for the same method name.

@IsTest
public class FlexibleMock implements System.StubProvider {
   private Map<String, Map<Integer, Object>> mockMethodAndArgumentResults;

   public FlexibleMock() {
       this.mockMethodAndArgumentResults = new Map<String, Map<Integer, Object>>();
   }

   public FlexibleMock mockMethodCall(String methodName, List<Object> argumentList, Object returnValue) {
       Integer hashCode = argumentList.hashCode();
       if (!this.mockMethodAndArgumentResults.containsKey(methodName)) {
           this.mockMethodAndArgumentResults.put(methodName, new Map<Integer, Object>{ hashCode => returnValue });
       } else {
           this.mockMethodAndArgumentResults.get(methodName).put(hashCode, returnValue);
       }
       return this;
   }

   public Object handleMethodCall( 
       Object stubbedObject,
       String stubbedMethodName,
       Type returnType,
       List<Type> listOfParamTypes,
       List<String> listOfParamNames,
       List<Object> listOfArgs
   ) {
       Integer argumentHash = listOfArgs != null ? listOfArgs.hashCode() : null;
       if (!this.mockMethodAndArgumentResults.containsKey(stubbedMethodName)) {
           throw new FlexibleMockException('Method mock results not registered: ' + stubbedMethodName);
       }
       Map<Integer, Object> methodMocks = this.mockMethodAndArgumentResults.get(stubbedMethodName);
       if (!methodMocks.containsKey(argumentHash)) {
           throw new FlexibleMockException(stubbedMethodName + ' mock results arguments not registered ' + argumentHash);
       }
       return methodMocks.get(argumentHash);
   }

   private class FlexibleMockException extends Exception {
   }
}

Limitations

So far this has worked well for me but it could definitely have problems. For instance using SObject arguments. Any difference in populated fields of the SObject instance will lead to a different hash code even though the field could be inconsequential to the method being mocked. We usually have full control over test data if we are mocking everything but it’s something to remember.

Of course the mocked classes have to be mockable in the first place. But that was not the point of this article. Neither is the way Mocks replace real objects. I focus on that a little bit in one of my other posts however.

Extra Curricular

I plucked out a test method from my current project (obscuring it ever so slightly) to add a bigger example of what the use of FlexibleMock can look like. It doesn’t matter what the test really does.

In general I need to prepare all my test records and then register mock versions of the classes with a Class Factory (mentioned already).

@IsTest
static void testSomeStuff() {
    Id ultimateParentId = TestUtilities.getFakeId(Account.getSObjectType());
    Account testAccount = childAccount(ultimateParentId);
    Account testChildAccount = childAccount(testAccount.Id, ultimateParentId);

    ADvendio__Commitment__c currentCommitment = currentCommitment(
        ultimateParentId,
        CommitmentParticipationService.CONTRACT_TYPE_COOPERATION
    );

    FlexibleMock mocks = new FlexibleMock()
        .mockMethodCall(
            'selectActiveCommitmentsByAccountId',
            new List<Object>{ new Set<Id>{ ultimateParentId } },
            new List<ADvendio__Commitment__c>{ currentCommitment }
        )
        .mockMethodCall(
            'findAllChildAccounts',
            new List<Object>{ new List<Account>{ testAccount } },
            new List<Account>{ testChildAccount }
        )
        .mockMethodCall('findAllChildAccounts', new List<Object>{ new List<Account>{ testChildAccount } }, new List<Account>())
        .mockMethodCall(
            'selectActiveCommitmentParticipationsByAccountId',
            new List<Object>{ new Set<Id>{ testAccount.Id, testChildAccount.Id } },
            new List<ADvendio__AccountCommitmentAssignment__c>()
        );
    ClassFactory.setMock(CommitmentSelector.class, (CommitmentSelector) Test.createStub(CommitmentSelector.class, mocks));
    ClassFactory.setMock(
        CommitmentParticipationSelector.class,
        (CommitmentParticipationSelector) Test.createStub(CommitmentParticipationSelector.class, mocks)
    );
    ClassFactory.setMock(AccountService.class, (AccountService) Test.createStub(AccountService.class, mocks));
...

This does not look very clean, does it. The FlexibleMock is not really at fault though. Just like the building of test records is delegated to some helper methods I can make this better by taking the mock “registration” out.

All the tests in the class do more or less the same thing. Which means they use the same Services and Selectors. Registering a selector that’s not strictly needed in one test, but is used in a bunch of others, is an acceptable redundancy for me here.

private static FlexibleMock getMockService() {
    FlexibleMock mocks = new FlexibleMock();
    ClassFactory.setMock(CommitmentSelector.class, (CommitmentSelector) Test.createStub(CommitmentSelector.class, mocks));
    ClassFactory.setMock(
        CommitmentParticipationSelector.class,
        (CommitmentParticipationSelector) Test.createStub(CommitmentParticipationSelector.class, mocks)
    );

    ClassFactory.setMock(AccountService.class, (AccountService) Test.createStub(AccountService.class, mocks)) 
    ClassFactory.setMock(
        MediaCampaignSelector.class,
        (MediaCampaignSelector) Test.createStub(MediaCampaignSelector.class, mocks)
    );
    return mocks;
}

It’s a bit cleaner this way. I only have to deal with the specific method calls I need in the given test. It’s still pretty wordy, but I found this to be quite enough. Depending on how many different test combinations there are and how similar their data, one could take this a bit further and externalise the full “setup”. But in my case this seemed unnecessary. This seems clean enough to me.

@IsTest
static void testSumeStuff() {
    Id ultimateParentId = TestUtilities.getFakeId(Account.getSObjectType());
    Account testAccount = childAccount(ultimateParentId);
    Account testChildAccount = childAccount(testAccount.Id, ultimateParentId);

    ADvendio__Commitment__c currentCommitment = currentCommitment(
        ultimateParentId,
        CommitmentParticipationService.CONTRACT_TYPE_COOPERATION
    );

    getMockService()
        .mockMethodCall(
            'selectActiveCommitmentsByAccountId',
            new List<Object>{ new Set<Id>{ ultimateParentId } },
            new List<ADvendio__Commitment__c>{ currentCommitment }
        )
        .mockMethodCall(
            'findAllChildAccounts',
            new List<Object>{ new List<Account>{ testAccount } },
            new List<Account>{ testChildAccount }
        )
        .mockMethodCall('findAllChildAccounts', new List<Object>{ new List<Account>{ testChildAccount } }, new List<Account>())
        .mockMethodCall(
            'selectActiveCommitmentParticipationsByAccountId',
            new List<Object>{ new Set<Id>{ testAccount.Id, testChildAccount.Id } },
            new List<ADvendio__AccountCommitmentAssignment__c>()
        );