[More Robust Mocking in Apex] A dynamic HTTP mock registry and a configurable Stub class to simplify and enhance mocking for Apex unit tests #salesforce #apex

HttpMockRegistry.cls

/**
 * A registry built around the Salesforce Mocking API that allows declarative mocking of HTTP callouts. Mocks responses
 * can be registered either for a specific endpoint and path or for all paths on an endpoint, with the former taking
 * precedence.
 */
@IsTest
public class HttpMockRegistry {

	// Default mock response for HTTP requests
	public static final HttpResponse DEFAULT_MOCK_RESPONSE = createSuccessResponse('Default mock response');

	// A registry of callout mocks, keyed by endpoint
	public static Map<String, CalloutMockConfig> calloutMocks { private get; private set; }

	// Static initialization
	static {
		calloutMocks = new Map<String, CalloutMockConfig>();
		Test.setMock(HttpCalloutMock.class, new CalloutResponder());
	}


	//==================================================================================================================
	// Mock Response context
	//==================================================================================================================

	/**
	 * Creates a mock REST context by setting the static RestContext's request and response fields.
	 */
	public static void mockRestContext() {
		mockRestContext(String.valueOf(Url.getOrgDomainUrl()), 'POST');
	}

	/**
	 * Creates a mock REST context by setting the static RestContext's request and response fields.
	 */
	public static void mockRestContext(String path, String httpMethod) {
		RestContext.request = new RestRequest();
		RestContext.request.resourcePath = path;
		RestContext.request.httpMethod = httpMethod;
	}


	//==================================================================================================================
	// Callout mocking
	//==================================================================================================================

	/**
	 * Mocks out a callout endpoint to return the default mock response for all requests.
	 * @param endpoint A callout endpoint
	 * endpoint and path
	 */
	public static void mockCallout(String endpoint) {
		getCalloutMockConfig(endpoint).setDefaultResponse(DEFAULT_MOCK_RESPONSE);
	}

	/**
	 * Mocks out a callout endpoint to return a given mock response by default for all requests.
	 * @param endpoint A callout endpoint
	 * @param response An HttpResponse to return whenever an HTTP request is sent to the provided
	 * endpoint and path
	 */
	public static void mockCallout(String endpoint, HttpResponse response) {
		getCalloutMockConfig(endpoint).setDefaultResponse(response);
	}

	/**
	 * Mocks out a callout endpoint to return a specified response for all requests to a particular path.
	 * @param endpoint A callout endpoint
	 * @param path The specific path on the given endpoint to mock
	 * @param response An HttpResponse to return whenever an HTTP request is sent to the provided
	 * endpoint and path
	 */
	public static void mockCallout(String endpoint, String path, HttpResponse mockResponse) {
		getCalloutMockConfig(endpoint).setResponse(path, mockResponse);
	}

	/**
	 * Mocks out a callout endpoint to return a specified responses for particular paths.
	 * @param endpoint A callout endpoint
	 * @param mockResponses A map of mock responses, keyed by the path for which they should
	 * be returned
	 */
	public static void mockCallout(String endpoint, Map<String, HttpResponse> mockResponses) {
		CalloutMockConfig mockConfig = getCalloutMockConfig(endpoint);
		for (String path : mockResponses.keySet()) {
			mockConfig.setResponse(path, mockResponses.get(path));
		}
	}

	/**
	 * Creates an HTTP success response with a given body.
	 * @param body Response body
	 * @return A 200 response with the given response body
	 */
	public static HttpResponse createSuccessResponse(String body) {
		return createResponse(200, 'OK', body);
	}

	/**
	 * Creates an HTTP response with a given status and body.
	 * @param statusCode Response status code
	 * @param status Response status message
	 * @param body Response body
	 * @return A response with given status and body
	 */
	public static HttpResponse createResponse(Integer statusCode, String status, String body) {
		HttpResponse response = new HttpResponse();
		response.setStatusCode(statusCode);
		response.setStatus(status);
		response.setBody(body);
		return response;
	}

	/**
	 * Returns the mock configuration for a given callout endpoint. If none has been established, it will be created
	 * and added to the configuration map.
	 * @param endpoint The endpoint for which to retrieve/create the mock configuration
	 */
	private static CalloutMockConfig getCalloutMockConfig(String endpoint) {
		CalloutMockConfig mockConfig = calloutMocks.get(endpoint);
		if (mockConfig == null) {
			mockConfig = new CalloutMockConfig();
			calloutMocks.put(endpoint, mockConfig);
		}
		return mockConfig;
	}


	//==================================================================================================================
	// Mock configs
	//==================================================================================================================

	/**
	 * A class to house the configuration of a mocked HTTP callout.
	 */
	private class CalloutMockConfig {
		public Map<String, HttpResponse> responses { get; set; }
		public HttpResponse defaultResponse { get; set; }

		/**
		 * Constructor.
		 */
		public CalloutMockConfig() {
			this.responses = new Map<String, HttpResponse>();
			this.defaultResponse = DEFAULT_MOCK_RESPONSE;
		}

		/**
		 * Configures the response for a specific path.
		 * @param path An HTTP service request path
		 * @param An HTTP response to return for the given path
		 */
		public void setResponse(String path, HttpResponse response) {
			responses.put(path, response);
		}

		/**
		 * Configures the default response for all paths.
		 * @param An HTTP response to return for all paths that aren't explicetly configured
		 */
		public void setDefaultResponse(HttpResponse defaultResponse) {
			this.defaultResponse = defaultResponse;
		}

		/**
		 * Returns a mock response for the given path.
		 * @param path An HTTP service request path
		 * @return A mock HTTP response for the given path
		 */
		public HttpResponse getResponse(String path) {
			HttpResponse response = responses.get(path);
			return (response != null) ? response : defaultResponse;
		}
	}


	//==================================================================================================================
	// HttpCalloutMock responder
	//==================================================================================================================

	/**
	 * An implementation of HttpCalloutMock used to repond to all HTTP requests.
	 */
	private class CalloutResponder implements HttpCalloutMock {
		public HttpResponse respond(HttpRequest request) {
			// Split request endpoint into base and path
			String base = request.getEndpoint().substringBefore('/');
			String path = request.getEndpoint().substringAfter('/');

			System.debug(String.format(
				'Mocking response for callout endpoint \'\'{0}\'\' on path \'\'{1}\'\'',
				new String[] { base, path }
			));

			// Get mock configuration for endpoint
			// If no mock response registered
			CalloutMockConfig mockConfig = calloutMocks.get(base);
			if (mockConfig == null) {
				throw new UnmockedCalloutException(String.format(
					'No mock response registered for callout to endpoint \'\'{0}\'\'',
					new String[] { base }
				));
			}

			// Find appropriate response in registry
			HttpResponse response = mockConfig.getResponse(path);
			System.debug('Response: ' + response);
			return response;
		}
	}

	//==================================================================================================================
	// Exceptions
	//==================================================================================================================

	// Exception thrown when a callout is intercepted that has not been properly mocked
	public class UnmockedCalloutException extends Exception {}
}

Stub.cls

/**
 * A configurable class/sObject instance stub that utilizes the Salesforce Stub API to return values from a method or
 * throw a specific exception.
 *
 * Because the Stub API requires one method to handle all invocations and the Reflection API does not provide a means of
 * invoking methods by name, it is not possible to stub out individual methods.
 *
 * A stub will throw a NoReturnValueException from all non-void methods not configured explicitly to avoid unexpected
 * NullPointerExceptions.
 */
@IsTest
public class Stub implements StubProvider {

	// Default mock exception
	public static final MockException MOCK_EXCEPTION = new MockException('Mock Exception');


	public Type type { get; private set; }
	public Object instance { get; private set; }
	public Set<String> invokedMethods { private get; private set; }
	public Map<String, Exception> exceptions { private get; private set; }
	public Map<String, Object> returnValues { private get; private set; }

	/**
	 * Constructor.
	 * @param typeToMock The class/SObject type that will be mocked by this stub
	 */
	public Stub(Type typeToMock) {
		this.type = typeToMock;
		this.instance = Test.createStub(typeToMock, this);
		this.invokedMethods = new Set<String>();
		this.exceptions = new Map<String, Exception>();
		this.returnValues = new Map<String, Object>();
	}


	//==================================================================================================================
	// Stub Provider
	//==================================================================================================================

	/**
	 * Handles a stubbed method call, per the StubProvider interface, by throwing an exception or returning a value, depending
	 * on what was configured for the method. If both a return value and exception have been set, the exception takes precedence.
	 * @param stubbedObject The stubbed object on which a method was invoked
	 * @param stubbedMethodName The name of the method that was invoked
	 * @param returnType The return type of the method that was invoked
	 * @param listOfParamTypes An ordered list of parameter types on the method that was invoked
	 * @param listOfParamNames An ordered list of parameter names on the method that was invoked
	 * @param listOfArgs An ordered list of arguments passed to the method
	 */
	public Object handleMethodCall(Object stubbedObject, String stubbedMethodName, Type returnType, List<Type> listOfParamTypes, List<String> listOfParamNames, List<Object> listOfArgs) {
		this.invokedMethods.add(stubbedMethodName);

		// If an exception has been configured for the invoked method, throw it
		// otherwise, if a return value has been configured for the method, return it
		if (exceptions.containsKey(stubbedMethodName)) {
			throw exceptions.get(stubbedMethodName);
		} else if (returnValues.containsKey(stubbedMethodName)) {
			return returnValues.get(stubbedMethodName);
		}

		// No result has been configured for the invoked method, throw an exception
		throw new NoReturnValueException(String.format(
			'No return value or exception has been configured for method \'\'{0}\'\' on stubbed class \'\'{1}\'\'',
			new String[] { stubbedMethodName }
		));
	}


	//==================================================================================================================
	// Configuration
	//==================================================================================================================

	/**
	 * Sets the value to be returned from a specific stubbed method.
	 * @param methodName The name of a method
	 * @param returnValue The value to return from the specified method when invoked
	 */
	public void setReturnValue(String methodName, Object returnValue) {
		returnValues.put(methodName, returnValue);
	}

	/**
	 * Configures a specified method to throw the default mock exception.
	 * @param methodName The name of a method
	 */
	public void setException(String methodName) {
		setException(methodName, MOCK_EXCEPTION);
	}

	/**
	 * Sets the exception to be thrown from a specific stubbed method.
	 * @param methodName The name of a method
	 * @param exceptionToThrow The exception to be thrown from the specified method when invoked
	 */
	public void setException(String methodName, Exception exceptionToThrow) {
		exceptions.put(methodName, exceptionToThrow);
	}


	//==================================================================================================================
	// Assertions
	//==================================================================================================================

	/**
	 * Asserts that a method with the given name has been invoked.
	 * @param methodName The name of the method in question
	 */
	public void assertInvoked(String methodName) {
		if (!invokedMethods.contains(methodName)) {
			Throw new MethodNotInvokedException(String.format(
				'Method {0}.{1}() not invoked',
				new String[] { type.getName(), methodName }
			));
		}
	}

	/**
	 * Asserts that a method with the given name has not been invoked.
	 * @param methodName The name of the method in question
	 */
	public void assertNotInvoked(String methodName) {
		if (invokedMethods.contains(methodName)) {
			Throw new MethodInvokedException(String.format(
				'Method {0}.{1}() invoked',
				new String[] { type.getName(), methodName }
			));
		}
	}


	//==================================================================================================================
	// Exceptions
	//==================================================================================================================

	// Default mock exception type
	public class MockException extends Exception {}

	public class NoReturnValueException extends Exception {}

	// Exception thrown when a stub's invocation assertion fails
	public class MethodInvokedException extends Exception {}

	// Exception thrown when a callout is intercepted that has not been properly mocked
	public class MethodNotInvokedException extends Exception {}
}