Add waits for flakey tests

Hi @PatrickRice, not sure if these waitfors are foolproof, as I couldn't reliably get the tests to flake locally, but I added waits. Let me know what you think, as I'd like to avoid the back-and-forth on MRs when there is flakiness. Thank you.