mymockito.java

package cz.cuni.mff.mymockito;

import net.sf.cglib.proxy.Enhancer;
import net.sf.cglib.proxy.MethodInterceptor;
import net.sf.cglib.proxy.MethodProxy;

import java.lang.reflect.Method;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Objects;

public class MyMockito {

    static Interceptor interceptorOfLastUsedMockedObject = null;


    /**
     * NEVER EVER USE IN ASYNC CONTEXT OR WITH MULTIPLE THREADS. IT WILL FAIL HARD!
     * @param objectClass
     * @param <T>
     * @return
     */
    @SuppressWarnings("unchecked")
    public static <T> T mock(Class<T> objectClass) {
        try {
            Enhancer enhancer = new Enhancer();
            enhancer.setSuperclass(objectClass); // parameter of mock method
            enhancer.setCallback(new Interceptor<T>(objectClass.getConstructor().newInstance()));
            return (T) enhancer.create(); // T is template arg of mock method
        } catch (Throwable e) {
            throw new RuntimeException(e);
        }
    }

    public static MockBuilder when(Object object) {
        return new MockBuilder(interceptorOfLastUsedMockedObject);

    }

    public static final class MockBuilder {
        Interceptor interceptor;
        MethodCall lastMethodCall;

        public MockBuilder(Interceptor interceptor) {
            this.interceptor = interceptor;
            lastMethodCall = interceptor.lastMethodCall;
        }

        public <T> void thenReturn(T o) {
            interceptor.registeredMethods.put(lastMethodCall, new MethodCallResult(o));
        }

        public <T> void thenThrows(Throwable throwable) {
            interceptor.registeredMethods.put(lastMethodCall, new MethodCallResult(throwable, false));
        }
    }

    private static final class MethodCallResult {
        Object result = null;
        Throwable t = null;

        public MethodCallResult(Object result) {
            this.result = result;
        }

        public MethodCallResult(Throwable t, boolean throwable) {
            this.t = t;
        }

        public boolean shouldThrow() {
            return t != null;
        }

        public Object getResult() {
            return result;
        }

        public Throwable getT() {
            return t;
        }
    }

    private static final class MethodCall {
        Method method;
        Object[] arguments;

        public MethodCall(Method method, Object[] arguments) {
            this.method = method;
            this.arguments = arguments;
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;
            MethodCall that = (MethodCall) o;
            return method.equals(that.method) &&
                    Arrays.equals(arguments, that.arguments);
        }

        @Override
        public int hashCode() {
            int result = Objects.hash(method);
            result = 31 * result + Arrays.hashCode(arguments);
            return result;
        }
    }


    private static final class Interceptor<T> implements MethodInterceptor {
        Object mockedObject;
        MethodCall lastMethodCall = null;
        HashMap<MethodCall, MethodCallResult> registeredMethods = new HashMap<>();
        public Interceptor(Object mockedObject) {
            this.mockedObject = mockedObject;
        }

        @Override
        public Object intercept(Object o, Method method, Object[] objects, MethodProxy methodProxy) throws Throwable {
            interceptorOfLastUsedMockedObject = this;

            MethodCall mc = new MethodCall(method, objects);
            lastMethodCall = mc;

            if (registeredMethods.containsKey(mc)) {
                MethodCallResult mcr = registeredMethods.get(mc);
                if (mcr.shouldThrow())
                    throw mcr.getT();
                else
                    return mcr.getResult();
            } else {
                return method.invoke(mockedObject, objects);
            }
        }
    }


}