SPARK-2277: refine UT