Minor unit test fix