Minor refactoring on unit test.