Implement awaiting action promises

src/core/timeoutScheduler.ts

@@ -25,21 +25,24 @@ export class TimeoutScheduler implements Disposable {
     this.actionWatchers.forEach(watcher => clearTimeout(watcher));
   }
 
-  setTimeout(action: () => void) {
-    if (this.counterExpirationMs)
-      this.resetCounterExpiration();
-
-    const timeoutIndex = Math.min(this.counter, this.timeouts.length - 1);
-    const timeout = this.timeouts[timeoutIndex];
-
-    const watcherId = setTimeout(() => {
-      this.actionWatchers.delete(watcherId);
-      clearTimeout(watcherId);
-      action();
-    }, timeout);
-    this.actionWatchers.add(watcherId);
-
-    this.counter++;
+  setTimeout(action: () => void | Promise<void>): Promise<void> {
+    return new Promise(resolve => {
+      if (this.counterExpirationMs)
+        this.resetCounterExpiration();
+
+      const timeoutIndex = Math.min(this.counter, this.timeouts.length - 1);
+      const timeout = this.timeouts[timeoutIndex];
+
+      const watcherId = setTimeout(async () => {
+        this.actionWatchers.delete(watcherId);
+        clearTimeout(watcherId);
+        await action();
+        resolve();
+      }, timeout);
+      this.actionWatchers.add(watcherId);
+
+      this.counter++;
+    });
   }
 
   resetCounter() {

tests/core/timeoutScheduler.test.ts

@@ -14,77 +14,100 @@ describe('TimeoutScheduler', () => {
 
   test('runs scheduled actions and clears timer', async () => {
     let counter = 0;
-    const action = () => counter++;
+    const action = () => { counter++; };
 
     scheduler = new TimeoutScheduler([500, 700]);
 
-    expect(scheduler.counter).toEqual(0);
+    expect(scheduler.counter).toBe(0);
     scheduler.setTimeout(action);
     await wait(500);
-    expect(counter).toEqual(1);
-    expect(scheduler.counter).toEqual(1);
+    expect(counter).toBe(1);
+    expect(scheduler.counter).toBe(1);
 
     scheduler.setTimeout(action);
     await wait(500);
-    expect(counter).toEqual(1);
+    expect(counter).toBe(1);
     await wait(200);
-    expect(counter).toEqual(2);
-    expect(scheduler.counter).toEqual(2);
+    expect(counter).toBe(2);
+    expect(scheduler.counter).toBe(2);
 
     scheduler.setTimeout(action);
     await wait(500);
-    expect(counter).toEqual(2);
+    expect(counter).toBe(2);
     await wait(200);
-    expect(counter).toEqual(3);
-    expect(scheduler.counter).toEqual(3);
+    expect(counter).toBe(3);
+    expect(scheduler.counter).toBe(3);
 
     scheduler.resetCounter();
-    expect(scheduler.counter).toEqual(0);
+    expect(scheduler.counter).toBe(0);
     scheduler.setTimeout(action);
     await wait(500);
-    expect(counter).toEqual(4);
-    expect(scheduler.counter).toEqual(1);
+    expect(counter).toBe(4);
+    expect(scheduler.counter).toBe(1);
   });
 
   test('runs scheduled actions and clears timer automatically', async () => {
     let counter = 0;
-    const action = () => counter++;
+    const action = () => { counter++; };
 
     scheduler = new TimeoutScheduler([500, 700], 1000);
 
-    expect(scheduler.counter).toEqual(0);
+    expect(scheduler.counter).toBe(0);
     scheduler.setTimeout(action);
     await wait(500);
-    expect(counter).toEqual(1);
-    expect(scheduler.counter).toEqual(1);
+    expect(counter).toBe(1);
+    expect(scheduler.counter).toBe(1);
 
     scheduler.setTimeout(action);
     await wait(500);
-    expect(counter).toEqual(1);
+    expect(counter).toBe(1);
     await wait(200);
-    expect(counter).toEqual(2);
-    expect(scheduler.counter).toEqual(2);
+    expect(counter).toBe(2);
+    expect(scheduler.counter).toBe(2);
 
     scheduler.setTimeout(action);
     await wait(500);
-    expect(counter).toEqual(2);
+    expect(counter).toBe(2);
     await wait(200);
-    expect(counter).toEqual(3);
-    expect(scheduler.counter).toEqual(3);
+    expect(counter).toBe(3);
+    expect(scheduler.counter).toBe(3);
 
     //Wait internal counter expiration
     await wait(1000);
-    expect(scheduler.counter).toEqual(0);
+    expect(scheduler.counter).toBe(0);
     scheduler.setTimeout(action);
     await wait(500);
-    expect(counter).toEqual(4);
-    expect(scheduler.counter).toEqual(1);
+    expect(counter).toBe(4);
+    expect(scheduler.counter).toBe(1);
+  });
+
+  test('waits for action promises', async () => {
+    jest.useFakeTimers({ doNotFake: ['setTimeout'] });
+
+    let counter = 0;
+    const action = (async () => {
+      await wait(100);
+      counter++;
+    });
+    scheduler = new TimeoutScheduler([300]);
+
+    await scheduler.setTimeout(action);
+    expect(counter).toBe(1);
+    expect(jest.getTimerCount()).toBe(0);
+
+    await scheduler.setTimeout(action);
+    expect(counter).toBe(2);
+    expect(jest.getTimerCount()).toBe(0);
+
+    await scheduler.setTimeout(action);
+    expect(counter).toBe(3);
+    expect(jest.getTimerCount()).toBe(0);
   });
 
   test('stops all scheduled tasks on dispose', () => {
     jest.useFakeTimers();
     const action = jest.fn();
-    scheduler = new TimeoutScheduler([10000, 10000]);
+    scheduler = new TimeoutScheduler([10000]);
 
     scheduler.setTimeout(action);
     scheduler.setTimeout(action);