Mock the Clock with an Aspect

Recently I have had problems with unit-tests that are time dependent. There are tests that test for time-out and similar time related issues. I felt that this mus be solvable by mocking the system clock. But I had problems doing it because I don’t want to change the whole implementation.

The best is if it is done by design so we have an interface like TimeProvider that gives us the current time, date and similar time related information. This can easily be mocked but the we have to use it every were with is not possible when some 3pp:s may depend on the time.

So how can we mock the time? In Java we can use an aspect to wrap around the System.currentTimeMillis() that is the time source for all time related in the platform (what I know of anyway). So we can use this to tick the time forward and easily use that to determine how much time has gone. But for this to work the code itself must be written with test in mind.

So I will provide I simple example of this here. The example is very constructed to just show the problem.

@RunWith(MockitoJUnitRunner.class)
public class TimingTest {

	@Mock
	DataGeneratorI dataGeneratorMock;

	@InjectMocks
	Timing timing;

	@Test(expected = RuntimeException.class)
	public void shouldTimoutAfterFourthCallAndThrowBadImplementation() {
		// Given
		given(dataGeneratorMock.generateDatasetWithIndex(anyInt())).will(
				new Answer<List<Integer>>() {

					@Override
					public List<Integer> answer(InvocationOnMock invocation)
							throws Throwable {
						// Simulate work
						Thread.sleep(10);
						return Arrays.asList(1);
					}
				});
		try {
			// When
			Date timeout = new Date(System.currentTimeMillis() + 35);
			timing.processLotsOfDataWithTimout(timeout);
		} finally {
			// Then
			verify(dataGeneratorMock, times(4)).generateDatasetWithIndex(anyInt());
		}
	}
}

public class Timing {

	private DataGeneratorI dataGeneretor;

	public void processLotsOfDataWithTimout(Date timeout) {
		Map<Integer, List<Integer>> generatedData = new HashMap<>();
		for (int i = 0; i < 10; ++i) {
			generatedData.put(i, dataGeneretor.generateDatasetWithIndex(i));
			checkForTimeout(timeout);
		}

	}

	private void checkForTimeout(Date timeout) {
		if (new Date().after(timeout)) {
			throw new RuntimeException("Timeout");
		}
	}
}

The problem here is the marked line in the test code. To simulate work the developer has chosen to sleep the current thread. If the test runs on a host that have high cpu load for the moment it can fail and make the test render an error even though the code works as it should. We can make the test more time independent by choosing higher values for the time-out and sleep time, but that will also make the test go slower. So a better an sustainable solution is to mock the system clock.

So lets create an aspect that we can use to achieve that.

@Aspect
public class SystemTimeMockAspect {

	@Pointcut("call(public long java.lang.System.currentTimeMillis())")
	void currentTimeMillis() {
	}
	
	@Around("currentTimeMillis()")
	public Object aroundSystemcurrentTimeMillis(ProceedingJoinPoint pjp) throws Throwable {	
		if (TimeProvider.instance().isTimedMocked()) {
			return TimeProvider.instance().currentTimeMillis();
		}
		return pjp.proceed();
	}
}

<aspectj>
	<aspects>
		<aspect name="se.aidium.mock.aspect.SystemTimeMockAspect" />
	</aspects>
	<weaver options="-verbose -showWeaveInfo" />
</aspectj>

The aop.xml is the glue between the aspect and the code. To make it work we need to be running java with the -javaagent:/path/to/aspectjweaver.jar option. I’m using load time weaving here so it’s not possible to weave the java rt classes (it’s possible with compile-time weaving). So we need to make one adjustment to the Timing class and that is to change the new Date().after(timout) to new Date(System.currentTimeMillis()).after(timout).

The advantages of mocking the clock is that now the test is independent of the real time and also we can make the test execute faster when we don’t have to wait for some time just to trigger the time-out.

So lets create the time provider and add a new test that demostrate the usage of this.

public class TimeProvider {
	private final static TimeProvider instance = new TimeProvider();

	public static TimeProvider instance() {
		return instance;
	}

	private long currentTimeMillis;
	private boolean mockTime;

	public boolean isTimedMocked() {
		return mockTime;
	}

	public void useMockTime(long currentTime) {
		currentTimeMillis = currentTime;
		mockTime = true;
	}

	public void useSystemTime() {
		mockTime = false;
	}

	public long currentTimeMillis() {
		return currentTimeMillis;
	}

	public void tick(int tick) {
		currentTimeMillis += tick;
	}
}

@Test(expected = RuntimeException.class)
public void shouldTimoutAfterFourthCallAndThrow() {
	// Given
	TimeProvider.instance().useMockTime(System.currentTimeMillis());
	given(dataGeneratorMock.generateDatasetWithIndex(anyInt())).will(new Answer<List<Integer>>() {

		@Override
		public List<Integer> answer(InvocationOnMock invocation) throws Throwable {
			// Simulate work
			TimeProvider.instance().tick(1);
			return Arrays.asList(1);
		}
	});

	try {
		// When
		Date timeout = new Date(System.currentTimeMillis() + 3);
		timing.processLotsOfDataWithTimout(timeout);
	} finally {
		// Then
		verify(dataGeneratorMock, times(4)).generateDatasetWithIndex(anyInt());
		TimeProvider.instance().useSystemTime();
	}
}

Compile-time weaving is an option if many 3pp:s uses time related classes that the aspect can’t weave during load-time and those 3pp:s can’t be wrapped by our own classes. That can be achieved with a weaving goal in eg maven or any other build system that supports weaving.

All code can be found at github.