Today I was writing unit test for a piece of code which required me to mock a iterator. My requirement was that I wanted to return true first time when hasNext() method and return false in the second iteration. It took me sometime to figure out how to write such a test case. My class for which I was writing test looks like as shown below.

import java.util.Iterator;

public class ResultFetcher {

	private ResultStore store;
	private ResultProcessor processor;

	public void fetchResults() {
		Iterator<String> iterator = store.resultIterator();
		while (iterator.hasNext()) {
			String result = iterator.next();
			processor.process(result);
		}
	}
}

The unit test code is shown below. The important line in the test is Mockito.when(iterator.hasNext()).thenReturn(true,false); which will return true first time and false second time.

import java.util.Iterator;

import org.junit.Test;
import org.mockito.Mockito;

public class ResultFetcherTest {

@Test
public void testFetchResults() {
ResultFetcher fetcher = new ResultFetcher();

ResultStore store = Mockito.mock(ResultStore.class);
ResultProcessor processor = Mockito.mock(ResultProcessor.class);
fetcher.store = store;
fetcher.processor = processor;

Iteratoriterator = Mockito.mock(Iterator.class);
Mockito.when(store.resultIterator()).thenReturn(iterator);
Mockito.when(iterator.hasNext()).thenReturn(true,false);
Mockito.when(iterator.next()).thenReturn("Hello");
fetcher.fetchResults();
Mockito.verify(processor).process("Hello");
}

}