UniversalMocks.java

/************************************************************

*** @author: Suraj Pillai
*** @description: A universal class for mocking in tests. Contains a method for setting the return value for any method. Another method returns the number of times a method was called

*/
@isTest
public with sharing class UniversalMocks implements System.StubProvider {
    public static Map<String, Object> returnValuesMap = new Map<String, Object>();
    public static Map<String, Integer> callCountsMap = new Map<String, Integer>();
    public static Map<String, List<Map<String, Object>>> argumentsMap = new Map<String, List<Map<String, Object>>>();

    private static String getClassNameFromStubbedObjectName(Object stubbedObject) {
        return String.valueOf(stubbedObject).split(':')[0].split('__')[0];
    }

    public static void setMock(Type mockedClass, String stubbedMethodName, Type returnType, Object returnValue) {
        String className = mockedClass.getName();
        String key = className + '||' + stubbedMethodName + '||' + returnType.getName();
        returnValuesMap.put(key, returnValue);
        if (!callCountsMap.containsKey(key))
            callCountsMap.put(key, 0);
    }

    //Can probably get rid of methods that don't accept a Type argument
    public static void setMock(String stubbedMethodName, Type returnType, Object returnValue) {
        String key = stubbedMethodName + '||' + returnType.getName();
        returnValuesMap.put(key, returnValue);
        if (!callCountsMap.containsKey(key))
            callCountsMap.put(key, 0);
    }

    public Object handleMethodCall(
        Object stubbedObject,
        String stubbedMethodName,
        Type returnType,
        List<Type> listOfParamTypes,
        List<String> listOfParamNames,
        List<Object> listOfArgs
    ) {
        String className = getClassNameFromStubbedObjectName(stubbedObject);
        String keyWithClassName = className + '||' + stubbedMethodName + '||' + returnType.getName();
        String keyWithoutClassName = stubbedMethodName + '||' + returnType.getName();
        String keyInUse = keyWithClassName;
        Integer count = callCountsMap.get(keyWithClassName);
        if (count == null) {
            count = callCountsMap.get(keyWithoutClassName);
            keyInUse = keyWithoutClassName;
            if (count == null) {
                count = 0;
            }
        }

        Map<String, Object> currentArgsMap = new Map<String, Object>();
        if (!argumentsMap.containsKey(keyInUse)) {
            argumentsMap.put(keyInUse, new List<Map<String, Object>>{ currentArgsMap });
        } else {
            argumentsMap.get(keyInUse).add(currentArgsMap);
        }

        for (Integer i = 0; i < listOfParamNames.size(); i++) {
            currentArgsMap.put(listOfParamNames[i], listOfArgs[i]);
        }
        callCountsMap.put(keyInUse, count + 1);
        Object returnValue = returnValuesMap.get(keyInUse);
        if (returnValue instanceof Exception) {
            throw (Exception) returnValue;
        }
        return returnValue;
    }

    public static Integer getCallCount(String stubbedMethodName, Type returnType) {
        return callCountsMap.get(stubbedMethodName + '||' + returnType.getName());
    }

    public static Integer getCallCount(Type theType, String stubbedMethodName, Type returnType) {
        String className = theType.getName();
        String keyWithClassName = className + '||' + stubbedMethodName + '||' + returnType.getName();
        return callCountsMap.get(keyWithClassName);
    }

    public static Object getMethodArgument(String stubbedMethodName, Type returnType, String argName, Integer invocationNumber) {
        String keyWithoutClassName = stubbedMethodName + '||' + returnType.getName();
        return argumentsMap.get(keyWithoutClassName).get(invocationNumber).get(argName);
    }

    public static Object getMethodArgument(Type theType, String stubbedMethodName, Type returnType, String argName, Integer invocationNumber) {
        String className = theType.getName();
        String keyWithClassName = className + '||' + stubbedMethodName + '||' + returnType.getName();
        return argumentsMap.get(keyWithClassName).get(invocationNumber).get(argName);
    }
}

Sample usage

UniversalMocks.setMock(DMLService.class, 'setSObjectList', IDMLService.class, dmlSvcInstance);
DMLService dmlSvcInstance = (DMLService) Test.createStub(DMLService.class, new UniversalMocks());

// run tests
...

// get the sobject list that was passed in for insert
        Submission__c[] submissionData = (Submission__c[]) UniversalMocks.getMethodArgument(
            DMLService.class,
            'setSObjectList',
            IDMLService.class,
            'sObjectList',
            2
        );
        system.assertEquals(submissionData[0].OwnerId, mockUser.Id, 'Record ownership not assigned to matching community user');