fix(testing/time): fix FakeTime.restoreFor accuracy for sync callbacks (#3531)

testing/time.ts

@@ -252,21 +252,26 @@ export class FakeTime {
   /**
    * Restores real time temporarily until callback returns and resolves.
    */
-  static async restoreFor<T>(
+  static restoreFor<T>(
     // deno-lint-ignore no-explicit-any
     callback: (...args: any[]) => Promise<T> | T,
     // deno-lint-ignore no-explicit-any
     ...args: any[]
   ): Promise<T> {
-    if (!time) throw new TimeError("no fake time");
-    let result: T;
+    if (!time) return Promise.reject(new TimeError("no fake time"));
     restoreGlobals();
     try {
-      result = await callback.apply(null, args);
-    } finally {
+      const result = callback.apply(null, args);
+      if (result instanceof Promise) {
+        return result.finally(() => overrideGlobals());
+      } else {
+        overrideGlobals();
+        return Promise.resolve(result);
+      }
+    } catch (e) {
       overrideGlobals();
+      return Promise.reject(e);
     }
-    return result;
   }
 
   /**

testing/time_test.ts

@@ -7,7 +7,7 @@ import {
   assertRejects,
   assertStrictEquals,
 } from "./asserts.ts";
-import { FakeTime } from "./time.ts";
+import { FakeTime, TimeError } from "./time.ts";
 import { _internals } from "./_time.ts";
 import { assertSpyCall, spy, SpyCall } from "./mock.ts";
 
@@ -304,6 +304,80 @@ Deno.test("FakeTime restoreFor restores real time temporarily", async () => {
   }
 });
 
+Deno.test("FakeTime restoreFor restores real time and re-overridden atomically", async () => {
+  const time: FakeTime = new FakeTime();
+  const fakeSetTimeout = setTimeout;
+  const actualSetTimeouts: (typeof setTimeout)[] = [];
+
+  try {
+    const asyncFn = async () => {
+      actualSetTimeouts.push(setTimeout);
+      await Promise.resolve();
+      actualSetTimeouts.push(setTimeout);
+      await Promise.resolve();
+      actualSetTimeouts.push(setTimeout);
+    };
+    const promise = asyncFn();
+    await new Promise((resolve) => {
+      FakeTime.restoreFor(() => setTimeout(resolve, 0));
+    });
+    await promise;
+    assertEquals(actualSetTimeouts, [
+      fakeSetTimeout,
+      fakeSetTimeout,
+      fakeSetTimeout,
+    ]);
+  } finally {
+    time.restore();
+  }
+});
+
+Deno.test("FakeTime restoreFor returns promise that resolved to result of callback", async () => {
+  const time: FakeTime = new FakeTime();
+
+  try {
+    const resultSync = await FakeTime.restoreFor(() => "a");
+    assertEquals(resultSync, "a");
+    const resultAsync = await FakeTime.restoreFor(() => Promise.resolve("b"));
+    assertEquals(resultAsync, "b");
+  } finally {
+    time.restore();
+  }
+});
+
+Deno.test("FakeTime restoreFor returns promise that rejected to error in callback", async () => {
+  const time: FakeTime = new FakeTime();
+
+  try {
+    await assertRejects(
+      () =>
+        FakeTime.restoreFor(() => {
+          throw new Error("Error in sync callback");
+        }),
+      Error,
+      "Error in sync callback",
+    );
+    await assertRejects(
+      () =>
+        FakeTime.restoreFor(() => {
+          return Promise.reject(new Error("Error in async callback"));
+        }),
+      Error,
+      "Error in async callback",
+    );
+  } finally {
+    time.restore();
+  }
+});
+
+Deno.test("FakeTime restoreFor returns promise that rejected to TimeError if FakeTime is uninitialized", async () => {
+  await assertRejects(
+    () => FakeTime.restoreFor(() => {}),
+    TimeError,
+    "no fake time",
+  );
+});
+
 Deno.test("delay uses real time", async () => {
   const time: FakeTime = new FakeTime();
   const start: number = Date.now();