Add back waitForStateTransitionToEnd to fix flakiness.

Bug: 202735477
Test: presubmit
Change-Id: I03c7e8b0c40177f4e6e38a7d9aa3e01b029effd6
2 files changed