Search code examples
javareflectionjava-8interfacemocking

Creating a mock library


I want to create a Mock Library class that implements InvocationHandler interface from Java Reflection.

This is the template I have created:

import java.lang.reflect.*;
import java.util.*;

class MyMock implements InvocationHandler {
        @Override
        public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
            // todo
        }
        
        public MyMock when(String method, Object[] args) {
            // todo
        }
        
        public void thenReturn(Object val) {
            // todo
        }
}

The when and thenReturn methods are chained methods.

Then when method registers the given mock parameters.

thenReturn method registers the expected return values for the given mock parameters.

Also, I want to throw java.lang.IllegalArgumentException if the proxied interface calls methods or uses parameters that are not registered.

This is a sample interface:

interface CalcInterface {
    int add(int a, int b);
    String add(String a, String b);
    String getValue();
}

Here we have two overloaded add methods.

This is a program to test the mock class I wanted to implement.

class TestApplication {     
        public static void main(String[] args) {
            MyMock m = new MyMock();
            CalcInterface ref = (CalcInterface) Proxy.newProxyInstance(MyMock.class.getClassLoader(), new Class[]{CalcInterface.class}, m);
            
            m.when("add", new Object[]{1,2}).thenReturn(3);
            m.when("add", new Object[]{"x","y"}).thenReturn("xy");
            
            System.out.println(ref.add(1,2)); // prints 3
            System.out.println(ref.add("x","y")); // prints "xy"
        }
}

This is the code which I have implemented so far to check the methods in CalcInterface:

class MyMock implements InvocationHandler {
        @Override
        public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
            int n = args.length;
            if(n == 2 && method.getName().equals("add")) {
                Object o1 = args[0], o2 = args[1];
                if((o1 instanceof String) && (o2 instanceof String)) {
                    String s1 = (String) o1, s2 = (String) o2;
                    return s1+ s2;
                } else if((o1 instanceof Integer) && (o2 instanceof Integer)) {
                    int s1 = (Integer) o1, s2 = (Integer) o2;
                    return s1+ s2;
                }
            }
            throw new IllegalArgumentException();
        }
        
        public MyMock when(String method, Object[] args) {
            return this;
        }
        
        public void thenReturn(Object val) {
        
        }
}

Here I am checking only for methods with the name add and having 2 arguments, with their type as String or Integer.

But I wanted to create this MyMock class in a general fashion, supporting different interfaces not just CalcInterface, and also supporting different methods not just the add method I implemented here.


Solution

  • You have to separate the builder logic from the object to build. The method when has to return something which remembers the arguments, so that the invocation of thenReturn still knows the context.

    For example

    public class MyMock implements InvocationHandler {
        record Key(String name, List<?> arguments) {
            Key { // stream().toList() creates an immutable list allowing null
                arguments = arguments.stream().toList();
            }
            Key(String name, Object... arg) {
                this(name, arg == null? List.of(): Arrays.stream(arg).toList());
            }
        }
        final Map<Key, Function<Object[], Object>> rules = new HashMap<>();
    
        @Override
        public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
            var rule = rules.get(new Key(method.getName(), args));
            if(rule == null) throw new IllegalStateException("No matching rule");
            return rule.apply(args);
        }
        public record Rule(MyMock mock, Key key) {
            public void thenReturn(Object val) {
                var existing = mock.rules.putIfAbsent(key, arg -> val);
                if(existing != null) throw new IllegalStateException("Rule already exist");
            }
            public void then(Function<Object[], Object> f) {
                var existing = mock.rules.putIfAbsent(key, Objects.requireNonNull(f));
                if(existing != null) throw new IllegalStateException("Rule already exist");
            }
        }
        public Rule when(String method, Object... args) {
            Key key = new Key(method, args);
            if(rules.containsKey(key)) throw new IllegalStateException("Rule already exist");
            return new Rule(this, key);
        }
    }
    

    This is already capable of executing your example literally, but also supports something like

    MyMock m = new MyMock();
    CalcInterface ref = (CalcInterface) Proxy.newProxyInstance(
            CalcInterface.class.getClassLoader(), new Class[]{CalcInterface.class}, m);
    
    m.when("add", 1,2).thenReturn(3);
    m.when("add", "x","y").thenReturn("xy");
    AtomicInteger count = new AtomicInteger();
    m.when("getValue").then(arg -> "getValue invoked " + count.incrementAndGet() + " times");
    
    System.out.println(ref.add(1,2)); // prints 3
    System.out.println(ref.add("x","y")); // prints "xy"
    System.out.println(ref.getValue()); // prints getValue invoked 1 times
    System.out.println(ref.getValue()); // prints getValue invoked 2 times
    

    Note that when you want to add support for rules beyond simple value matching, a hash lookup will not work anymore. In that case you have to resort to a data structure you have to search linearly for a match.

    The example above uses newer Java features like record classes but it shouldn’t be too hard to rewrite it for previous Java versions if required.


    It’s also possible to redesign this code to use the real builder pattern, i.e. to use a builder to describe the configuration prior to creating the actual handler/mock instance. This allows the handler/mock to use an immutable state:

    public class MyMock2 {
        public static Builder builder() {
            return new Builder();
        }
        public interface Rule {
            Builder thenReturn(Object val);
            Builder then(Function<Object[], Object> f);
        }
        public static class Builder {
            final Map<Key, Function<Object[], Object>> rules = new HashMap<>();
    
            public Rule when(String method, Object... args) {
                Key key = new Key(method, args);
                if(rules.containsKey(key))
                    throw new IllegalStateException("Rule already exist");
                return new RuleImpl(this, key);
            }
            public <T> T build(Class<T> type) {
                Map<Key, Function<Object[], Object>> rules = Map.copyOf(this.rules);
                return type.cast(Proxy.newProxyInstance(type.getClassLoader(),
                    new Class[]{ type }, (proxy, method, args) -> {
                       var rule = rules.get(new Key(method.getName(), args));
                       if(rule == null) throw new IllegalStateException("No matching rule");
                       return rule.apply(args);
                    }));
    
            }
        }
        record RuleImpl(MyMock2.Builder builder, Key key) implements Rule {
            public Builder thenReturn(Object val) {
                var existing = builder.rules.putIfAbsent(key, arg -> val);
                if(existing != null) throw new IllegalStateException("Rule already exist");
                return builder;
            }
            public Builder then(Function<Object[], Object> f) {
                var existing = builder.rules.putIfAbsent(key, Objects.requireNonNull(f));
                if(existing != null) throw new IllegalStateException("Rule already exist");
                return builder;
            }
        }
        record Key(String name, List<?> arguments) {
            Key { // stream().toList() createns an immutable list allowing null
                arguments = arguments.stream().toList();
            }
            Key(String name, Object... arg) {
                this(name, arg == null? List.of(): Arrays.stream(arg).toList());
            }
        }
    }
    

    which can be used like

    AtomicInteger count = new AtomicInteger();
    CalcInterface ref = MyMock2.builder()
            .when("add", 1,2).thenReturn(3)
            .when("add", "x","y").thenReturn("xy")
            .when("getValue")
                .then(arg -> "getValue invoked " + count.incrementAndGet() + " times")
            .build(CalcInterface.class);
    
    System.out.println(ref.add(1,2)); // prints 3
    System.out.println(ref.add("x","y")); // prints "xy"
    System.out.println(ref.getValue()); // prints getValue invoked 1 times
    System.out.println(ref.getValue()); // prints getValue invoked 2 times