Mockito的使用及原理浅析

Mockito的使用及原理浅析

在微服务横行的年代，一个服务可能依赖了若干个其它服务，而这些被依赖的服务，又极可能依赖了别的服务，于是构成了一个复杂的依赖链。而日常开发中，经常会苦恼于某某上游接口没数据而所使得UT达不到预期的目标，所兴现在已经有了多种Mock Toolkit，很好的解决了这一痛点。比较常用的有EasyMock,JMockit,Mockito,PowerMock等，而本文主要介绍Mockito的简单使用，并对其实现原理进行简单分析

Mockito的使用

通常情况下，我们会在UT里使用@Mock、@Spy以及@InjectMocks等注解来Mock掉一些不可控的上下游接口，为了简化说明，我们将例子简单化。
假设有这样的一个服务

public class ServiceA {

    @Autowired
    private ServiceB serviceB;
    
    public String hello(String name) {
        return "ServiceA.hello:" + name;
    }
    
    public String helloB(String name) {
        return serviceB.hello(name);
    }
}

@Mock

现在我们来将它mock掉

public class MockTest {
    
    @Mock
    private ServiceA seriveA;

    @Before
    public void init() {
        MockitoAnnotations.initMocks(this);
    }
    
    @Test
    public void testA() {
        Mockito.when(seriveA.hello("aaa")).thenReturn("111").thenReturn("222");  //打桩
        System.out.println(seriveA.hello("aaa"));   //111                        //真实调用
        System.out.println(seriveA.hello("aaa"));   //222
        System.out.println(seriveA.hello("bbb"));   //null
    }
}

为了便于描述，先简单说明一下：上面serviceA.hello(“aaa”)调用结果做为入参传递给when()方法，我们称为打桩，而后面的调用是我们称为真实调用。
Mockito提供了通俗易用的API，@Mock注释用于对目标对象进行mock，生成立mock对象，when…thenReturn…用于模拟返回结果，当调用某某方法时返回我们想要的预定结果， thenThrow则用于模拟目标方法抛出异常，而我们也可以使用方法链模拟返回多个mock结果，在多次调用目标方法时，Mockito会依次为我们返回预先设定的mock结果。

@InjectMocks

ServiceA依赖了ServiceB，如果我们要mock的不是ServiceA，而是ServiceB，那我们可以在UT里，给ServiceB加注释@Mock，给ServiceA加上注释@InjectMocks，这时Mockito会自动为ServiceA注入mock对象ServiceB

public class MockTest2 {
    
    @InjectMocks
    private ServiceA serviceA;
    @Mock
    private ServiceB serviceB;
    
    @Before
    public void init() {
        MockitoAnnotations.initMocks(this);
    }
    
    @Test
    public void testB() {
        Mockito.when(serviceB.hello("aaa")).thenReturn("111").thenReturn("222");
        System.out.println(serviceA.helloB("aaa"));   //111
        System.out.println(serviceA.helloB("aaa"));   //222
        System.out.println(serviceA.helloB("bbb"));   //null
    }
}

从上面的用法中，整个mock的过程大概可以分为三步：
1）@Mock注解声明了要将对象进行mock
2）使用MockitoAnnotations.initMocks告诉Mockito要对当前对象进行mock处理。当然你用可以使用@RunWith(MockitoJUnitRunner.class)，殊途同归。
3）使用when…thenReturn来模拟返回
接下来我们来分析一下其实现原理。

Mockito的实现原理

mock对象的创建

首先我们看看MockitoAnnotations.initMocks到底做了什么？一起进入源码。

public static void initMocks(Object testClass) {
        if (testClass == null) {
            throw new MockitoException("testClass cannot be null. For info how to use @Mock annotations see examples in javadoc for MockitoAnnotations class");
        }

        AnnotationEngine annotationEngine = new GlobalConfiguration().getAnnotationEngine();
        Class<?> clazz = testClass.getClass();

        //below can be removed later, when we get read rid of deprecated stuff
        if (annotationEngine.getClass() != new DefaultMockitoConfiguration().getAnnotationEngine().getClass()) {
            //this means user has his own annotation engine and we have to respect that.
            //we will do annotation processing the old way so that we are backwards compatible
            while (clazz != Object.class) {
                scanDeprecatedWay(annotationEngine, testClass, clazz);
                clazz = clazz.getSuperclass();
            }
        }

        //anyway act 'the new' way
        annotationEngine.process(testClass.getClass(), testClass);
    }

这里有个AnnotationEngine，从名称中我们大概可以猜到Mockito是要对Annotation比如：@Mock,@Spy,@InjectMocks等进行处理，以此创建出Mock对象。
AnnotationEngine有三个实现，DefaultAnnotationEngine,SpyAnnotationEngine,InjectingAnnotationEngine,这三个AnnotationEngine合分别处理@Mock、@Spy和InjectMocks注解。这里选择DefaultAnnotationEngine进行分析。

public void process(Class<?> clazz, Object testInstance) {
        Field[] fields = clazz.getDeclaredFields();
        for (Field field : fields) {
            boolean alreadyAssigned = false;
            for(Annotation annotation : field.getAnnotations()) {           
                Object mock = createMockFor(annotation, field);
                if (mock != null) {
                    throwIfAlreadyAssigned(field, alreadyAssigned);                    
                    alreadyAssigned = true;                    
                    try {
                        new FieldSetter(testInstance, field).set(mock);
                    } catch (Exception e) {
                        throw new MockitoException("Problems setting field " + field.getName() + " annotated with "
                                + annotation, e);
                    }
                }        
            }
        }
    }

在这里看到createMockFor，由于代码还比较长，略去代码，直接看下面的流程图

从流程是我们大概知道，MockitoAnnotations.initMocks(this)的过程就是创建mock对象的过程，Mockito在1.x版本中是采用cglib动态创建被mock的目标类的代理类(2.x版本改用了bytebuddy)，代理类是被mock目标类的子类，因此Mocktio无法mock被定义为final的类和方法。而在创建目标mock对象时，Mocktio利用cglib给mock对象加上了方法拦截器(MethodInterceptorFilter)，这样当mock对象的方法被调用时，会被拦截器进行拦截处理。其中关键的代理类的生成和对象实例是创建是在ClassImposterizer中完成的，代码如下：

public Class<Factory> createProxyClass(Class<?> mockedType, Class<?>... interfaces) {
        if (mockedType == Object.class) {
            mockedType = ClassWithSuperclassToWorkAroundCglibBug.class;
        }
        
        Enhancer enhancer = new Enhancer() {
            @Override
            @SuppressWarnings("unchecked")
            protected void filterConstructors(Class sc, List constructors) {
                // Don't filter
            }
        };
        Class<?>[] allMockedTypes = prepend(mockedType, interfaces);
		enhancer.setClassLoader(SearchingClassLoader.combineLoadersOf(allMockedTypes));
        enhancer.setUseFactory(true);
        if (mockedType.isInterface()) {
            enhancer.setSuperclass(Object.class);
            enhancer.setInterfaces(allMockedTypes);
        } else {
            enhancer.setSuperclass(mockedType);
            enhancer.setInterfaces(interfaces);
        }
        enhancer.setCallbackTypes(new Class[]{MethodInterceptor.class, NoOp.class});
        enhancer.setCallbackFilter(IGNORE_BRIDGE_METHODS);
        if (mockedType.getSigners() != null) {
            enhancer.setNamingPolicy(NAMING_POLICY_THAT_ALLOWS_IMPOSTERISATION_OF_CLASSES_IN_SIGNED_PACKAGES);
        } else {
            enhancer.setNamingPolicy(MockitoNamingPolicy.INSTANCE);
        }

        enhancer.setSerialVersionUID(42L);
        
        try {
            return enhancer.createClass(); 
        } catch (CodeGenerationException e) {
            if (Modifier.isPrivate(mockedType.getModifiers())) {
                throw new MockitoException("\n"
                        + "Mockito cannot mock this class: " + mockedType 
                        + ".\n"
                        + "Most likely it is a private class that is not visible by Mockito");
            }
            throw new MockitoException("\n"
                    + "Mockito cannot mock this class: " + mockedType 
                    + "\n" 
                    + "Mockito can only mock visible & non-final classes."
                    + "\n" 
                    + "If you're not sure why you're getting this error, please report to the mailing list.", e);
        }
    }
    
    private Object createProxy(Class<Factory> proxyClass, final MethodInterceptor interceptor) {
        Factory proxy;
        try {
            proxy = instantiator.newInstance(proxyClass);
        } catch (InstantationException e) {
            throw new MockitoException("Unable to create mock instance of type '" + proxyClass.getSuperclass().getSimpleName() + "'", e);
        }
        proxy.setCallbacks(new Callback[] {interceptor, SerializableNoOp.SERIALIZABLE_INSTANCE });
        return proxy;
    }

when().thenReturn()

分析完mock对象的构建，接下来说when()…thenReturn()…，其中when()方法入参是Object对象，我们一般是先调用目标方法然后将其返回值传给when()（该步骤也被称为打桩），而对目标方法的调用会先被拦截器(MethodInterceptorFilter)拦截，最后由MockHandlerImpl进行处理，代码如下：

public Object handle(Invocation invocation) throws Throwable {
		if (invocationContainerImpl.hasAnswersForStubbing()) {
            // stubbing voids with stubVoid() or doAnswer() style
            InvocationMatcher invocationMatcher = matchersBinder.bindMatchers(
                    mockingProgress.getArgumentMatcherStorage(),
                    invocation
            );
            invocationContainerImpl.setMethodForStubbing(invocationMatcher);
            return null;
        }
        VerificationMode verificationMode = mockingProgress.pullVerificationMode();

        InvocationMatcher invocationMatcher = matchersBinder.bindMatchers(
                mockingProgress.getArgumentMatcherStorage(),
                invocation
        );

        mockingProgress.validateState();

        // if verificationMode is not null then someone is doing verify()
        if (verificationMode != null) {
            // We need to check if verification was started on the correct mock
            // - see VerifyingWithAnExtraCallToADifferentMockTest (bug 138)
            if (((MockAwareVerificationMode) verificationMode).getMock() == invocation.getMock()) {
                VerificationDataImpl data = createVerificationData(invocationContainerImpl, invocationMatcher);
                verificationMode.verify(data);
                return null;
            } else {
                // this means there is an invocation on a different mock. Re-adding verification mode
                // - see VerifyingWithAnExtraCallToADifferentMockTest (bug 138)
                mockingProgress.verificationStarted(verificationMode);
            }
        }

        // prepare invocation for stubbing
        invocationContainerImpl.setInvocationForPotentialStubbing(invocationMatcher);
        OngoingStubbingImpl<T> ongoingStubbing = new OngoingStubbingImpl<T>(invocationContainerImpl);
        mockingProgress.reportOngoingStubbing(ongoingStubbing);

        // look for existing answer for this invocation
        StubbedInvocationMatcher stubbedInvocation = invocationContainerImpl.findAnswerFor(invocation);

        if (stubbedInvocation != null) {
            stubbedInvocation.captureArgumentsFrom(invocation);
            return stubbedInvocation.answer(invocation);
        } else {
             Object ret = mockSettings.getDefaultAnswer().answer(invocation);

            // redo setting invocation for potential stubbing in case of partial
            // mocks / spies.
            // Without it, the real method inside 'when' might have delegated
            // to other self method and overwrite the intended stubbed method
            // with a different one. The reset is required to avoid runtime exception that validates return type with stubbed method signature.
            invocationContainerImpl.resetInvocationForPotentialStubbing(invocationMatcher);
            return ret;
        }
	}

主要逻辑是判断当前的调用是打桩还是真实调用，如果是打桩，则将当前调用信息(Invocation)交给InvocationContainerImpl进行记录，如果不是，则从InvocationContainerImpl中查找StubbedInvocationMatcher，而StubbedInvocationMatcher维护着一个由mock方法的mock结果组成的队列，可以根据客户端的调用依次返回mock结果。
thenReturn()是OngoingStubbing接口定义的方法，由BaseStubbing抽象类实现，并最终会调用OngoingStubbingImpl的thenAnswer(),最后将mock交给InvocationContainerImpl进行记录。OngoingStubbingImpl部分源码如下：

public OngoingStubbing<T> thenAnswer(Answer<?> answer) {
        if(!invocationContainerImpl.hasInvocationForPotentialStubbing()) {
            new Reporter().incorrectUseOfApi();
        }

        invocationContainerImpl.addAnswer(answer);
        return new ConsecutiveStubbing<T>(invocationContainerImpl);
    }

整个when().thenReturn()的过程如下：

执行完when()…thenReturn…，则完成了打桩，接下来调用目标方法，MockHandlerImpl从InvoncationContainerImpl查询到mock结果，直接返回给客户端(TestCase)。
了解完整个mock的实现过程，可以自己实现一个简单的mock工具。

mock工具的简易实现

先来个总览

MockTool：主入口

public class MockTool {

    private static final MockMethodInterceptor INTERCEPTOR = new MockMethodInterceptor();
    
    public static <T> T mock(Class<T> type) {
        Enhancer enhancer = new Enhancer();
        enhancer.setSuperclass(type);
        enhancer.setCallbackFilter((method)->method.isBridge() ? 1 : 0);//忽略bridge方法
        enhancer.setCallbacks(new Callback[] {INTERCEPTOR, NoOp.INSTANCE});
        return type.cast(enhancer.create());
    }
	
	public static <T> Stubbing when(T call) {
		return new Stubbing();
	}
}

MockMethodInterceptor：拦截目标方法，记录并返回mock结果

public class MockMethodInterceptor implements MethodInterceptor {

    @Override
    public Object intercept(Object obj, Method method, Object[] args, MethodProxy proxy) throws Throwable {
        Invocation invocation = new Invocation(obj.getClass(), method, args);
        Object value = InvocationContainer.findAnswerFor(invocation);
        if (value == null) {
            //查询不到mock结果,视为要打桩
            InvocationContainer.setCurrentInvocation(invocation);
        }
        return value;
    }
    
}

InvocationContainer：记录当前要mock的方法，以及mock结果

public class InvocationContainer {

    private static Invocation currentInvocation;

    private static final ConcurrentHashMap<Invocation, Queue<Object>> INVOCATION_RETURN_MAP = new ConcurrentHashMap<>();

    public static void setCurrentInvocation(Invocation invocation) {
        InvocationContainer.currentInvocation = invocation;
    }

    public static Invocation getCurrentInvocation() {
        return currentInvocation;
    }

    public static void addAnswer(Invocation invocation, Object value) {
        Queue<Object> queue = INVOCATION_RETURN_MAP.get(invocation);
        if (queue == null) {
            queue = new ConcurrentLinkedQueue<>();
            INVOCATION_RETURN_MAP.put(invocation, queue);
        }
        queue.offer(value);
    }

    public static Object findAnswerFor(Invocation invocation) {
        Queue<Object> queue = INVOCATION_RETURN_MAP.get(invocation);
        if (queue == null) {
            return queue;
        }
        return queue.size() > 1 ? queue.poll() : queue.peek();
    }
}

Invocation：抽象对方法的调用

@Setter
@Getter
@AllArgsConstructor
@EqualsAndHashCode
public class Invocation {
    
    private Class<?> mockClass;
    private Method method;
    private Object[] args;
    
}

Stubbing：将mock结果写入mock结果容器

public class Stubbing {

    public Stubbing thenReturn(Object result) {
        if (InvocationContainer.getCurrentInvocation() != null) {
            InvocationContainer.addAnswer(InvocationContainer.getCurrentInvocation(), result);
        } else {
            throw new IllegalStateException("api使用姿势不对");
        }
        return this;
    }
}

最后走一个测试，基本达标。

public class TestMockTool {

    @Test
    public void testMock() {
        ServiceA serviceA = MockTool.mock(ServiceA.class);
        MockTool.when(serviceA.hello("aaa")).thenReturn("111").thenReturn("222");
        System.out.println(serviceA.hello("aaa")); //111
        System.out.println(serviceA.hello("aaa")); //222
        System.out.println(serviceA.hello("bbb")); //null
    }
}