This is not a full fix for the flakiness. One of the test cases will
still be flaky. But at least this fixes the other test cases thus
improving stability.