Fix race condition in a test.
Fix race condition in a test.

PiperOrigin-RevId: 398208848
Change-Id: I021fefc624df53741b4b5a365e15de58930fcac6

deepmind/reverb,deepmind/reverb,deepmind/reverb,deepmind/reverb