Fix a race condition making test flaky. PiperOrigin-RevId: 480104467 (cherry picked from commit cb3a6725)
