Minor changes related to unit test