[Service Bus] Fix flaky test - Peek until we get the message (Azure#13346)

Investigated at Azure#13336 to figure out and fix the flakyness.

Sometimes, the service returns zero messages for the peek call and is forcing us to retry the peek operation until we get the message so the test is green.

(We may have other tests that may fail with the same, we can extend this fix to those tests as needed.)

Author: HarshaNalluru

sdk/servicebus/service-bus/test/renewLock.spec.ts

@@ -7,7 +7,7 @@ const assert = chai.assert;
 import chaiAsPromised from "chai-as-promised";
 chai.use(chaiAsPromised);
 import { delay } from "rhea-promise";
-import { TestMessage } from "./utils/testUtils";
+import { checkWithTimeout, TestMessage } from "./utils/testUtils";
 import {
   ServiceBusClientForTests,
   createServiceBusClientForTests,
@@ -74,9 +74,21 @@ describe("Message Lock Renewal", () => {
     const testMessage = TestMessage.getSample();
     await sender.sendMessages(testMessage);
 
-    const [peekedMsg] = await receiver.peekMessages(1);
+    let peekedMsg: ServiceBusReceivedMessage | undefined;
+    should.equal(
+      await checkWithTimeout(
+        async () => {
+          [peekedMsg] = await receiver.peekMessages(1);
+          return peekedMsg != undefined;
+        },
+        100,
+        5000
+      ),
+      true,
+      "Unable to peek messages in 5000 milliseconds"
+    );
     try {
-      await receiver.renewMessageLock(peekedMsg);
+      await receiver.renewMessageLock(peekedMsg!);
       assert.fail("renewMessageLock should have failed");
     } catch (error) {
       should.equal(error.message, InvalidOperationForPeekedMessage);

sdk/servicebus/service-bus/test/utils/testUtils.ts

@@ -151,13 +151,13 @@ export enum TestClientType {
  * Keep checking whether the predicate is true after every `1000 ms`(default value is 1 second) (= delayBetweenRetriesInMilliseconds)
  */
 export async function checkWithTimeout(
-  predicate: () => boolean,
+  predicate: () => boolean | Promise<boolean>,
   delayBetweenRetriesInMilliseconds: number = 1000,
   maxWaitTimeInMilliseconds: number = 10000
 ): Promise<boolean> {
   const maxTime = Date.now() + maxWaitTimeInMilliseconds;
   while (Date.now() < maxTime) {
-    if (predicate()) return true;
+    if (await predicate()) return true;
     await delay(delayBetweenRetriesInMilliseconds);
   }
   return false;