fix(defer): recover from error (#3933)

* feat(defer): implement waitable callback

* fix(defer): recover from callback error

* test(defer): remove useless wait test

src/lib/utils/__tests__/defer.ts

@@ -9,7 +9,7 @@ describe('defer', () => {
 
     expect(fn).toHaveBeenCalledTimes(0);
 
-    await Promise.resolve();
+    await deferred.wait();
 
     expect(fn).toHaveBeenCalledTimes(1);
   });
@@ -24,7 +24,7 @@ describe('defer', () => {
 
     expect(fn).toHaveBeenCalledTimes(0);
 
-    await Promise.resolve();
+    await deferred.wait();
 
     expect(fn).toHaveBeenCalledTimes(1);
   });
@@ -39,15 +39,15 @@ describe('defer', () => {
 
     expect(fn).toHaveBeenCalledTimes(0);
 
-    await Promise.resolve();
+    await deferred.wait();
 
     expect(fn).toHaveBeenCalledTimes(1);
 
     deferred();
     deferred();
     deferred();
 
-    await Promise.resolve();
+    await deferred.wait();
 
     expect(fn).toHaveBeenCalledTimes(2);
   });
@@ -64,7 +64,7 @@ describe('defer', () => {
 
     expect(fn).toHaveBeenCalledTimes(0);
 
-    await Promise.resolve();
+    await deferred.wait();
 
     expect(fn).toHaveBeenCalledTimes(0);
   });
@@ -81,13 +81,13 @@ describe('defer', () => {
 
     expect(fn).toHaveBeenCalledTimes(0);
 
-    await Promise.resolve();
+    await deferred.wait();
 
     expect(fn).toHaveBeenCalledTimes(0);
 
     deferred();
 
-    await Promise.resolve();
+    await deferred.wait();
 
     expect(fn).toHaveBeenCalledTimes(1);
   });
@@ -104,8 +104,49 @@ describe('defer', () => {
 
     expect(fn).toHaveBeenCalledTimes(0);
 
-    await Promise.resolve();
+    await deferred.wait();
 
     expect(fn).toHaveBeenCalledTimes(1);
   });
+
+  it('throws an error when `wait` is called before the deferred function', () => {
+    const fn = jest.fn();
+    const deferred = defer(fn);
+
+    expect(() => deferred.wait()).toThrowErrorMatchingInlineSnapshot(
+      `"The deferred function should be called before calling \`wait()\`"`
+    );
+  });
+
+  it('recovers a deferred function that throws an error', async () => {
+    const fn = jest.fn();
+    const deferred = defer(fn);
+
+    fn.mockImplementation(() => {
+      throw new Error('FAIL');
+    });
+
+    deferred();
+
+    expect(fn).toHaveBeenCalledTimes(0);
+
+    try {
+      await deferred.wait();
+    } catch {
+      // The test verifies that the function is able to recover. We don't want
+      // to terminate the test on this expected error.
+    }
+
+    expect(fn).toHaveBeenCalledTimes(1);
+
+    fn.mockImplementation();
+
+    deferred();
+
+    expect(fn).toHaveBeenCalledTimes(1);
+
+    await deferred.wait();
+
+    expect(fn).toHaveBeenCalledTimes(2);
+  });
 });

src/lib/utils/defer.ts

@@ -1,35 +1,48 @@
 const nextMicroTask = Promise.resolve();
 
 type Callback = (...args: any[]) => void;
-type Cancellable = Callback & {
+type Defer = Callback & {
+  wait(): Promise<void>;
   cancel(): void;
 };
 
-const defer = (callback: Callback): Cancellable => {
+const defer = (callback: Callback): Defer => {
   let progress: Promise<void> | null = null;
   let cancelled = false;
 
-  const fn: Cancellable = (...args) => {
+  const fn: Defer = (...args) => {
     if (progress !== null) {
       return;
     }
 
     progress = nextMicroTask.then(() => {
+      progress = null;
+
       if (cancelled) {
-        progress = null;
         cancelled = false;
         return;
       }
 
       callback(...args);
-      progress = null;
     });
   };
 
+  fn.wait = () => {
+    if (progress === null) {
+      throw new Error(
+        'The deferred function should be called before calling `wait()`'
+      );
+    }
+
+    return progress;
+  };
+
   fn.cancel = () => {
-    if (progress !== null) {
-      cancelled = true;
+    if (progress === null) {
+      return;
     }
+
+    cancelled = true;
   };
 
   return fn;