Make identity tests run 230x faster (#5482)

Improve the runtime of Identity unit tests by avoiding waiting real-world clock time when testing code that uses delays, such as polling for token updates.

sdk/identity/identity/src/credentials/deviceCodeCredential.ts

@@ -2,10 +2,11 @@
 // Licensed under the MIT License.
 
 import qs from "qs";
-import { TokenCredential, GetTokenOptions, AccessToken, delay } from "@azure/core-http";
+import { TokenCredential, GetTokenOptions, AccessToken } from "@azure/core-http";
 import { IdentityClientOptions, IdentityClient, TokenResponse } from "../client/identityClient";
 import { AuthenticationError, AuthenticationErrorName } from "../client/errors";
 import { createSpan } from "../util/tracing";
+import { delay } from "../util/delay";
 import { CanonicalCode } from "@azure/core-tracing";
 
 /**

sdk/identity/identity/src/util/delay.ts

@@ -0,0 +1,25 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+let testFunction: ((t: number) => Promise<void>) | undefined;
+
+/**
+ * A wrapper for setTimeout that resolves a promise after t milliseconds.
+ * @internal
+ * @param {number} t The number of milliseconds to be delayed.
+ * @returns {Promise<void>} Resolved promise
+ */
+export function delay(t: number): Promise<void> {
+  if (testFunction) {
+    return testFunction(t);
+  }
+
+  return new Promise((resolve) => setTimeout(() => resolve(), t));
+}
+
+/**
+ * @internal
+ */
+export function _setDelayTestFunction(func?: (t: number) => Promise<void>): void {
+  testFunction = func;
+}

sdk/identity/identity/test/authTestUtils.ts

@@ -3,12 +3,12 @@
 
 import assert from "assert";
 import { IdentityClientOptions } from "../src";
+import { _setDelayTestFunction } from "../src/util/delay";
 import {
   HttpHeaders,
   HttpOperationResponse,
   WebResource,
   HttpClient,
-  delay,
   RestError
 } from "@azure/core-http";
 
@@ -66,7 +66,6 @@ export class MockAuthHttpClient implements HttpClient {
     this.requests.push(httpRequest);
 
     if (this.mockTimeout) {
-      await delay(httpRequest.timeout);
       throw new RestError("Request timed out", RestError.REQUEST_SEND_ERROR);
     }
 
@@ -125,3 +124,82 @@ export async function assertRejects(
     assert.ok(expected(error), message || "The error didn't pass the assertion predicate.");
   }
 }
+
+export function setDelayInstantlyCompletes(): void {
+  _setDelayTestFunction(() => Promise.resolve());
+}
+
+export interface DelayInfo {
+  resolve: () => void;
+  reject: (e: Error) => void;
+  promise: Promise<void>;
+  timeout: number;
+}
+
+export class DelayController {
+  private _waitPromise?: Promise<DelayInfo>;
+  private _resolve?: (info: DelayInfo) => void;
+  private _pendingRequests: DelayInfo[] = [];
+
+  private removeDelayInfo(info: DelayInfo): void {
+    const index = this._pendingRequests.indexOf(info);
+    if (index >= 0) {
+      this._pendingRequests.splice(index, 1);
+    }
+  }
+
+  delayRequested(timeout: number): Promise<void> {
+    let resolveFunc: () => void;
+    let rejectFunc: (e: Error) => void;
+    const promise = new Promise<void>((resolve, reject) => {
+      resolveFunc = resolve;
+      rejectFunc = reject;
+    });
+    const info: DelayInfo = {
+      resolve: resolveFunc!,
+      reject: rejectFunc!,
+      promise,
+      timeout
+    };
+    this._pendingRequests.push(info);
+
+    const removeThis = (): void => {
+      this.removeDelayInfo(info);
+    };
+
+    promise.then(removeThis).catch(removeThis);
+
+    if (this._resolve) {
+      this._resolve(info);
+      this._resolve = undefined;
+      this._waitPromise = undefined;
+    }
+
+    return promise;
+  }
+
+  getPendingRequests(): DelayInfo[] {
+    return this._pendingRequests;
+  }
+
+  waitForDelay(): Promise<DelayInfo> {
+    if (!this._waitPromise) {
+      this._waitPromise = new Promise<DelayInfo>((resolve) => {
+        this._resolve = resolve;
+      });
+    }
+    return this._waitPromise;
+  }
+}
+
+export function createDelayController(): DelayController {
+  const controller = new DelayController();
+  _setDelayTestFunction((t) => {
+    return controller.delayRequested(t);
+  });
+  return controller;
+}
+
+export function restoreDelayBehavior(): void {
+  _setDelayTestFunction();
+}

sdk/identity/identity/test/node/deviceCodeCredential.spec.ts

@@ -2,10 +2,16 @@
 // Licensed under the MIT License.
 
 import assert from "assert";
-import { delay } from "@azure/core-http";
 import { TestTracer, setTracer, SpanGraph } from "@azure/core-tracing";
 import { AbortController } from "@azure/abort-controller";
-import { MockAuthHttpClient, assertRejects } from "../authTestUtils";
+import {
+  MockAuthHttpClient,
+  assertRejects,
+  setDelayInstantlyCompletes,
+  restoreDelayBehavior,
+  createDelayController,
+  DelayController
+} from "../authTestUtils";
 import { AuthenticationError, ErrorResponse } from "../../src/client/errors";
 import {
   DeviceCodeCredential,
@@ -27,7 +33,12 @@ const pendingResponse: ErrorResponse = {
 };
 
 describe("DeviceCodeCredential", function() {
-  this.timeout(10000); // eslint-disable-line no-invalid-this
+  before(() => {
+    setDelayInstantlyCompletes();
+  });
+  after(() => {
+    restoreDelayBehavior();
+  });
 
   it("sends a device code request and returns a token when the user completes it", async function() {
     const mockHttpClient = new MockAuthHttpClient({
@@ -223,7 +234,13 @@ describe("DeviceCodeCredential", function() {
         { status: 200, parsedBody: deviceCodeResponse },
         { status: 400, parsedBody: pendingResponse },
         { status: 400, parsedBody: pendingResponse },
-        { status: 401, parsedBody: { error: "invalid_client", error_description: "The request body must contain..."} }
+        {
+          status: 401,
+          parsedBody: {
+            error: "invalid_client",
+            error_description: "The request body must contain..."
+          }
+        }
       ]
     });
 
@@ -242,32 +259,43 @@ describe("DeviceCodeCredential", function() {
     });
   });
 
-  it("cancels polling when abort signal is raised", async function() {
-    const mockHttpClient = new MockAuthHttpClient({
-      authResponse: [
-        { status: 200, parsedBody: deviceCodeResponse },
-        { status: 400, parsedBody: pendingResponse },
-        { status: 400, parsedBody: pendingResponse },
-        { status: 200, parsedBody: { access_token: "token", expires_in: 5 } }
-      ]
+  describe("tests with delays", function() {
+    let delayController: DelayController;
+    before(() => {
+      delayController = createDelayController();
     });
 
-    const credential = new DeviceCodeCredential(
-      "tenant",
-      "client",
-      (details) => assert.equal(details.message, deviceCodeResponse.message),
-      mockHttpClient.identityClientOptions
-    );
-
-    const abortController = new AbortController();
-    const getTokenPromise = credential.getToken("scope", { abortSignal: abortController.signal });
-    await delay(1500); // Long enough for device code request and one polling request
-    abortController.abort();
-
-    const token = await getTokenPromise;
+    it("cancels polling when abort signal is raised", async function() {
+      const mockHttpClient = new MockAuthHttpClient({
+        authResponse: [
+          { status: 200, parsedBody: deviceCodeResponse },
+          { status: 400, parsedBody: pendingResponse },
+          { status: 400, parsedBody: pendingResponse },
+          { status: 200, parsedBody: { access_token: "token", expires_in: 5 } }
+        ]
+      });
+
+      const credential = new DeviceCodeCredential(
+        "tenant",
+        "client",
+        (details) => assert.equal(details.message, deviceCodeResponse.message),
+        mockHttpClient.identityClientOptions
+      );
 
-    assert.strictEqual(token, null);
-    assert.strictEqual(mockHttpClient.requests.length, 2);
+      const abortController = new AbortController();
+      const getTokenPromise = credential.getToken("scope", { abortSignal: abortController.signal });
+      // getToken ends up calling pollForToken which normally has a 1000ms delay.
+      // This code allows us to control the delay programatically in the test.
+      let delay = await delayController.waitForDelay();
+      delay.resolve();
+      delay = await delayController.waitForDelay();
+      // abort the request before the second poll is allowed to complete
+      abortController.abort();
+      delay.resolve();
+      const token = await getTokenPromise;
+      assert.strictEqual(token, null);
+      assert.strictEqual(mockHttpClient.requests.length, 2);
+    });
   });
 
   it("sends a device code request and returns a token with tracing", async function() {

sdk/identity/identity/test/node/managedIdentityCredential.spec.ts

@@ -97,7 +97,7 @@ describe("ManagedIdentityCredential", function() {
 
     assert.strictEqual(authDetails.requests[0].timeout, 5000);
     assert.strictEqual(authDetails.token, null);
-  }).timeout(10000);
+  });
 
   it("doesn't try IMDS endpoint again once it can't be detected", async function() {
     const mockHttpClient = new MockAuthHttpClient({ mockTimeout: true });