Expectations and Enums for nicely readable tests

While working reasonably closely with some guys from Thoughtworks (at my last job), I picked up a rather nice trick for making tests readable. They (or at least the guys I was working with) have a tendency to write inner classes called Expectations (which often mature into standalone classes if you use them for integration tests). You set these up, define some state you want, then call a verify() or similarly named method, that checks that everything is, indeed, in the state you want. One of the nice things you can use these to do is define a "base" state, so you don't constantly find yourself writing assertions that test that the code you've written hasn't modified something it shouldn't have.

Clear as mud, right? Maybe a snippet will help. First, a little description of the problem I'll be using to illustrate things. We've been given the task of programming to an interface. We'll be providing an implementation of a message processor, and one of the requirements that the interface lays on us is that we count the number of events that happen. Here's the interface:

	package com.blogspot.swishbob;
	public interface MessageEventHandler {
	public abstract void onEvent(MessageEvent sourceEvent);
	public abstract long getAssignedMessagesCount();
	public abstract long getCompletedMessagesCount();
	public abstract long getConsumedMessagesCount();
	public abstract long getReassignedMessagesCount();
	public abstract long getTimedOutMessagesCount();
	}

view raw MessageEventHandler.java hosted with ❤ by GitHub

We're aiming to get to a nice set of tests that look a little like this:

	@Test
	public void shouldOnlyIncrementMessageReassignedCountWhenMessageReassignedEventReceived() throws Exception {
	givenSourceEvent(MESSAGE_REASSIGNED);
	eventDrivenObject.onEvent(sourceEvent);
	expect(MessageCounter.REASSIGNED, 1).only();
	}
	private void givenSourceEvent(MessageEventType type) {
	sourceEvent = new MessageEvent(type);
	}

view raw ExpectationTestSnippet.java hosted with ❤ by GitHub

There's a relatively obvious way to code your counted method, you have it directly call each of the individual get count methods. It's not a bad way, in fact in plenty of cases it's simpler than what I want to show you, and therefore better, but it's not what I want to show here.

I'm a big fan of Java enums. There's heaps of cool stuff you can do with them that people never think to use. I possibly overcompensate (there's very little you can do with an enum that you couldn't do with an ordinary class and a bunch of static instance variables held within it, for example), but one of the things I like to use them for is to eliminate case statements, particularly the sort which are used to decide which method to call and nothing much else. The sort you find in this sort of event driven programming all over the place ...

For reference, the MessageEvent and MessageEventType we're working with:

	package com.blogspot.swishbob;
	public class MessageEvent {
	private final MessageEventType type;
	public MessageEvent(MessageEventType type) {
	this.type = type;
	}
	public MessageEventType getType() {
	return type;
	}
	}

view raw MessageEvent.java hosted with ❤ by GitHub

	package com.blogspot.swishbob;
	public enum MessageEventType {
	MESSAGE_REASSIGNED,
	MESSAGE_ASSIGNED,
	MESSAGE_TIMED_OUT,
	MESSAGE_COMPLETED
	}

view raw MessageEventType.java hosted with ❤ by GitHub

The key to this trick is that we create an enum that knows how to get the required count from our MessageEventHandler. We do this by creating the enum with an abstract getCountFrom(MessageEventHandler eventHandler) method, which the individual counter will implement with the correct call. Wrap it in a nice descriptive assertEquals call, and we get something like this:

	package com.blogspot.swishbob;
	import static org.junit.Assert.assertEquals;
	enum MessageCounter {
	ASSIGNED {
	@Override
	protected long getCountFrom(MessageEventHandler eventHandler) {
	return eventHandler.getAssignedMessagesCount();
	}
	},
	COMPLETED {
	@Override
	protected long getCountFrom(MessageEventHandler eventHandler) {
	return eventHandler.getCompletedMessagesCount();
	}
	} ,
	CONSUMED {
	@Override
	protected long getCountFrom(MessageEventHandler eventHandler) {
	return eventHandler.getConsumedMessagesCount();
	}
	},
	REASSIGNED {
	@Override
	protected long getCountFrom(MessageEventHandler eventHandler) {
	return eventHandler.getReassignedMessagesCount();
	}
	},
	TIMED_OUT {
	@Override
	protected long getCountFrom(MessageEventHandler eventHandler) {
	return eventHandler.getTimedOutMessagesCount();
	}
	};
	public void verifyExpected(int count, MessageEventHandler eventHandler) {
	assertEquals("Count for event type " + this + " was wrong", count, getCountFrom(eventHandler));
	}
	protected abstract long getCountFrom(MessageEventHandler eventHandler);
	}

view raw MessageCounter.java hosted with ❤ by GitHub

The only thing missing from the picture now is the Expectation class itself. Which is pretty straightforward. The nice part here is that the Expectation itself is completely divorced from the different event types and how they're counted. If we add a new one, it doesn't need to be touched.

	package com.blogspot.swishbob;
	import java.util.HashMap;
	import java.util.Map;
	class Expectation {
	private final Map<MessageCounter, Integer> expectedCounts = new HashMap<MessageCounter, Integer>();
	private final MessageEventHandler eventHandler;
	public static Expectation expect(final MessageCounter assignedMessages, final int count) {
	return new Expectation(eventDrivenObject).and(assignedMessages, count);
	}
	private Expectation(MessageEventHandler eventHandler) {
	this.eventHandler = eventDrivenObject;
	for (MessageCounter type : MessageCounter.values()) {
	expectedCounts.put(type, 0);
	}
	}
	public Expectation and(final MessageCounter messageType, final int count) {
	expectedCounts.put(messageType, count);
	return this;
	}
	public void only() {
	for (MessageCounter type : expectedCounts.keySet()) {
	type.verifyExpected(expectedCounts.get(type), eventHandler);
	}
	}
	}

view raw Expectation.java hosted with ❤ by GitHub

And that's it, really. A neat trick which I think helps. Probably more useful when things start getting scarily complex than in a case this simple, but hopefully somebody will find it useful some day.