How do I return mocked item only on the x:th call?

Question

I have run into a problem using PowerMock where I feel forced to create ugly code if I want to return a mocked item only once. As an example I have the following code:

mockMap = spy(new HashMap());
HashMap<String, String> normalMap = new HashMap<>();
HashMap<String, String> normalMap2 = new HashMap<>();
HashMap<String, String> normalMap3 = new HashMap<>();
whenNew(HashMap.class).withNoArguments()
    .thenReturn(normalMap)
    .thenReturn(mockMap)
    .thenReturn(normalMap2)
    .thenReturn(normalMap3);

This of course works, but it feels very clunky, especially as I need to create a new hashmap for every new call.

So my question is: is there a way to tell PowerMock that it should stop interfering after a set amount of calls?

edit: After reading the answers, I got the following:

final AtomicInteger count = new AtomicInteger(0);
whenNew(HashMap.class).withNoArguments().thenAnswer(invocation -> {
switch (count.incrementAndGet())
{
        case 1:
            return mockMap;
        default:
            return new HashMap<String, String>();
    }
});

However it gives me a StackOverFlowError using the following code:

describe("test", () -> {
    beforeEach(() -> {
        mockStatic(HashMap.class);
        final AtomicInteger count = new AtomicInteger(0);
        whenNew(HashMap.class).withNoArguments().thenAnswer(invocation ->
        {
            switch (count.incrementAndGet())
            {
                case 5:
                    return mockMap;
                default:
                    return new HashMap<String, String>();
            }
        });
    });
    it("Crashes on getting a new HashMap", () -> {
        HashMap map = new HashMap<>();
        map.put("test", "test");
        HashMap normal = new HashMap<String, String>();
        expect(normal.containsValue("test")).toBeTrue();
    });
});

Worth noting is that I don't have a mockStatic(HashMap.class) in my larger tests that get the same errors(which I get rid of if I remove the mockStatic call)

A solution that works(but feels like a workaround) is building on that by editing the default statement into:

default:
    normalMap.clear();
    return normalMap;

Answer 1

You could use thenReturn with multiple arguments like this:

whenNew(HashMap.class).withNoArguments()
    .thenReturn(normalMap, mockMap, normalMap2, normalMap3);

Or write your own Answer like this:

whenNew(HashMap.withNoArguments()).doAnswer(new Answer() {
    private int count = 0;

    public Object answer(InvocationOnMock invocation) {
        // implement your logic
        // if (count ==0) etc.
    }
});

Answer 2

You can use Mockito's org.mockito.stubbing.Answer:

final AtomicInteger count = new AtomicInteger(0);
PowerMockito.whenNew(HashMap.class).withNoArguments()
    .thenAnswer(new Answer<HashMap<String, String>>() {

    @Override
    public HashMap<String, String> answer(InvocationOnMock invocation) throws Throwable {
        count.incrementAndGet();
        switch (count.get()) {
            case 1: // first call, return normalMap
                return normalMap;
            case 2: // second call, return mockMap
                return mockMap;
            case 3: // third call, return normalMap2
                return normalMap2;
            default: // forth call (and all calls after that), return normalMap3
                return normalMap3;
        }
    }
});

Note that I had to use java.util.concurrent.atomic.AtomicInteger and declare it as a final variable for 2 reasons:

• variables used inside the anonymous class must be final (because Java defined it that way)
• if I create a final int, I can't do count++

Note: in this solution, you must also change all normalMap's to be final

---

Actually, if all your normalMap's are the same, you can just do:

PowerMockito.whenNew(HashMap.class).withNoArguments()
    .thenAnswer(new Answer<HashMap<String, String>>() {

    @Override
    public HashMap<String, String> answer(InvocationOnMock invocation) throws Throwable {
        count.incrementAndGet();
        if (count.get() == 2) { // second call
            return mockMap;
        }
        return normalMap; // don't forget to make normalMap "final"
        // or if you prefer: return new HashMap<String, String>();
    }
});

---

PS: as stated in this answer, instead of using AtomicInteger, you can also create an int counter inside the anonymous class (it's up to you to choose, as both work):

PowerMockito.whenNew(HashMap.class).withNoArguments()
    .thenAnswer(new Answer<HashMap<String, String>>() {

    private int count = 0;

    @Override
    public HashMap<String, String> answer(InvocationOnMock invocation) throws Throwable {
        count++;
        if (count == 2) { // second call
            return mockMap;
        }

        return normalMap; // don't forget to make normalMap "final"
        // or if you prefer: return new HashMap<String, String>();
    }
});

And this also works with the switch solution as well.