Iterators are your friends!

We all know that Iterators nicely decouple the container from the implementation details of traversing the container. Today I was optimizing an existing piece of code and had to work through the details this post talks about.

More about Iterators here. I’m going to talk about  a specific advantage of using iterators – improved testability of the code.

Let’s consider this simple class AwesomeParser. It keeps an InputStream and reads tokens from the stream. The stream has a header which contains the number of tokens in the stream.

class AwesomeParser {
    private final InputStream fileStream;
    private int numberOfTokens;
    private int tokensRead;

    public AwesomeParser(fileInputStream) {
        fileStream = fileInputStream;       
        numberOfTokens = 0;
        tokensRead = 0;
    }

    public void readFileHeader() {
        numberOfTokens = extractTokensCountFromFileHeader();
    }

    public Token read() {
        Token token = parseTokenFromFile();
        tokensRead++;
        return token;
    }

    public boolean hasMoreTokens() {
        return tokensRead < numberOfTokens;
    }
}

A sample consumer of this class would look something like,

class Consumer {
    private  final AwesomeParser parser;
    
    public Consumer(AwesomeParser parser) {
        this.parser = parser;
    }

    public void fooBar() {
        AwesomeParser parser = new AwesomeParser(fooStream);
        parser.readFileHeader();
        while (parser.hasMoreTokens()) {
            parser.read();
        }
    }
}

Looks like a nice abstraction. Now let’s see how the unit tests for our consumer would look like.

• We mock out the parser.
• Now the only way to test fooBar() is to make hasMoreToken() return different values for different calls. For e.g,

when(parser.hasMoreTokens()).thenReturn(true).thenReturn(true).thenReturn(false);
when (parser.read()).thenReturn(dummyToken);

Wait a second, what would happen if I modify fooBar() to something like,

public void fooBar() {
    AwesomeParser parser = new AwesomeParser(fooStream);
    parser.readFileHeader();
    if (parser.hasMoreTokens()) {
        log(“Woohoo work to do”);
    }
    while (parser.hasMoreTokens()) {
        parser.read();
    }
}

We have broken our assumptions about hasMoreTokens() in fooBar(). The unit tests will start failing one after another as they are so brittle. So, what do we do now? Iterators to the rescue :)

Let’s redesign our AwesomeParser.

class AwesomeParser implements Iterable {
    private final InputStream fileStream;
    private int numberOfTokens;
    private int tokensRead;

    public AwesomeParser(fileInputStream) {
        fileStream = fileInputStream;       
        numberOfTokens = 0;
        tokensRead = 0;
    }

    public void readFileHeader() {
        numberOfTokens = extractTokensCountFromFileHeader();
    }

    private Token read() {
        Token token = parseTokenFromFile();
        tokensRead++;
        return token;
    }

    private boolean hasMoreTokens() {
        return tokensRead < numberOfTokens;
    }

    /*  Warning: Usually you don't keep the traversal state of a container
     *  outside the iterator. The code will miserably fail if some one
     *  calls iterator() on AwesomeParser multiple times, since we have 
     *  one single underlying stream. You can always create a deep copy of
     *  it in the constructor of AwesomeParserIterator.
     */
    public Iterator iterator() {
        return new AwesomeParserIterator();
    }
    
    class AwesomeParserIterator implements Iterator {

        @Override
        public boolean hasNext() {
            return hasMoreTokens();
        }

        @Override
        public Token next() {
            return read();
        }

        @Override
        public void remove() {
            // never mind
        }
    }
}

Let’s try this abstraction. If we have to unit test the consumer of our new design,

List <Token> tokens = new ArrayList<>();
//populate tokens
Iterator <Token> mockIterator = tokens.iterator();
when (AwesomeParser.iterator()).thenReturn(mockIterator);

Essentially we have covered the underlying parser’s stream with a simple dummy list of tokens and faked the iterator returned by our AwesomeParser with an iterator to our dummy list.

What have we achieved?

The unit tests are no longer brittle and would just test the behavior of the consumer of AwesomeParser and have no assumption about how we traverse. Neat isn’t it?

Ciao,
Surya.