Mock Same Method with Different Parameters

1. Overview

When mocking a method in Java, it can be useful to receive different responses based on the parameters passed in. In this article, we’ll look at different ways to achieve that goal depending on the complexity of our requirements.

2. Setup

First, let’s create an example service which we want to mock:

class ExampleService {
    int getValue(int arg){
        return 1;
    }
}

We’ve got a very simple service with a single method. The method has a single int as an argument and returns an int. Note that the argument and the return value have no relation, so by default, it’ll always return 1.

3. Limitations of Consecutive Stubbing

Let’s look at consecutive stubbing and what we can and can’t do with it. We can use consecutive stubbing to get different arguments from our mock in order regardless of the input we provide. This obviously lacks control over matching particular inputs to desired outputs but is useful in many cases. To do this we pass the method we want to stub to when(). We then chain a call to thenReturn() providing the responses in the order we want them:

@Test
void givenAMethod_whenUsingConsecutiveStubbing_thenExpectResultsInOrder(){
    when(exampleService.getValue(anyInt())).thenReturn(9, 18, 27);
    assertEquals(9, exampleService.getValue(1));
    assertEquals(18, exampleService.getValue(1));
    assertEquals(27, exampleService.getValue(1));
    assertEquals(27, exampleService.getValue(1));
}

We can see from the assertions that despite always passing in 1 as the parameter, we received back the expected values in order. Once all the values have been returned, all future calls will return the final value, as seen in the fourth call in our test.

4. Stubbing Calls for Different Parameters

We can extend our use of when() and thenReturn() to return different values for different parameters: 

@Test
void givenAMethod_whenStubbingForMultipleArguments_thenExpectDifferentResults() {
    when(exampleService.getValue(10)).thenReturn(100);
    when(exampleService.getValue(20)).thenReturn(200);
    when(exampleService.getValue(30)).thenReturn(300);

    assertEquals(100, exampleService.getValue(10));
    assertEquals(200, exampleService.getValue(20));
    assertEquals(300, exampleService.getValue(30));
}

The argument for when() is the method we want to stub, along with the value we want to specify a response for. By chaining the call to when() with thenReturn(), we’ve instructed the mock to return the requested value when the correct argument is received. We’re free to apply as many of these as we want to our mock to handle a range of inputs. We’ll receive the requested return value every time the expected input value is provided.

5. Using thenAnswer()

A more complex option, offering maximum control, is to use thenAnswer(). This allows us to take the arguments, perform any computation on them we want, and then return a value that’ll be outputted when interacting with the mock:

@Test
void givenAMethod_whenUsingThenAnswer_thenExpectDifferentResults() {
    when(exampleService.getValue(anyInt())).thenAnswer(invocation -> {
        int argument = (int) invocation.getArguments()[0];
        int result;
        switch (argument) {
        case 25:
            result = 125;
            break;
        case 50:
            result = 150;
            break;
        case 75:
            result = 175;
            break;
        default:
            result = 0;
        }
        return result;
    });
    assertEquals(125, exampleService.getValue(25));
    assertEquals(150, exampleService.getValue(50));
    assertEquals(175, exampleService.getValue(75));
}

Above, we’ve grabbed the arguments using getArguments() on the provided invocation object. We’ve assumed a single int argument here, but we could’ve catered for several different types. We also could’ve checked that there was at least one argument and the cast to an int was successful. To demonstrate the capabilities, we’ve used a switch statement to return different values based on the input. At the bottom, we can see from the assertions that our mocked service returns the results of the switch statement.

This option allows us to handle an unlimited amount of inputs with a single call to when(). The sacrifice is the readability and maintainability of the tests.

6. Conclusion

In this tutorial, we’ve seen three ways of configuring a mocked method to return different values. We looked at consecutive stubbing and saw that it’s useful for returning known values in order for any input but is very limited beyond that. Using when() combined with thenReturn() for each potential input offers a simple solution with improved control.  Alternatively, we can use thenAnswer() for maximum control over the relationship between the given input and the expected output. All three are useful depending on the requirements under test.