Fixing flaky unit test.

Reviewed at https://reviews.apache.org/r/20916/