[Yujun] Improved the test.