Slight tweak to unit test setup.