Add resume check to other WorkProfileTests to fix flakiness.

- Tests were flaky with waitForStateTransitionToEnd
- Tests stopped flaking with both waitForResume and waitForStateTransitionToEnd
- waitForResume may be enough to fix the flakiness, so
  removing waitForStateTransitionToEnd for now, if flakiness
  returns we can add it back

Bug: 202735477
Test: presubmit
Change-Id: Ib9cc16ea7815387582486517a87249d0053475f6