From 68cd2a3f47382a7983d9bcf18dff02c5139806b6 Mon Sep 17 00:00:00 2001
From: Erik Moura <erikian@erikian.dev>
Date: Mon, 24 Jul 2023 23:43:32 -0300
Subject: [PATCH] test: update the mock to handle exclusive locks

---
 tests/setup.ts | 141 ++++++++++++++++++++++++++++++++++++++++++++-----
 tests/utils.ts |  22 ++++++++
 2 files changed, 149 insertions(+), 14 deletions(-)

diff --git a/tests/setup.ts b/tests/setup.ts
index 3f4d79b13a..83e006644a 100644
--- a/tests/setup.ts
+++ b/tests/setup.ts
@@ -4,6 +4,7 @@ import { createSerializer } from 'enzyme-to-json';
 import { configure as mobxConfigure } from 'mobx';
 
 import { ElectronFiddleMock } from './mocks/mocks';
+import { getOrCreateMapValue } from './utils';
 
 enzymeConfigure({ adapter: new Adapter() });
 
@@ -75,47 +76,159 @@ delete (window as any).localStorage;
 window.navigator = window.navigator ?? {};
 (window.navigator.clipboard as any) = {};
 
+type MockLock = Lock & {
+  abortController: AbortController;
+};
+
 class FakeNavigatorLocks implements LockManager {
   locks = {
-    held: new Set<Lock>(),
-    pending: new Set<Lock>(),
+    held: new Map<string, Map<LockMode, Set<MockLock>>>(),
   };
 
   query = async () => {
     const result = {
-      held: [...this.locks.held],
-      pending: [...this.locks.pending],
+      held: [...this.locks.held.values()].reduce((acc, item) => {
+        acc.push(...[...item.get('exclusive')!.values()]);
+        acc.push(...[...item.get('shared')!.values()]);
+
+        return acc;
+      }, [] as MockLock[]),
     };
 
     return result as LockManagerSnapshot;
   };
 
-  /**
-   * WIP. Right now, this is a **very** naive mock that will just happily grant a shared lock when one is requested,
-   * but I'll add some bookkeeping and expand it to cover the exclusive lock case as well.
-   *
-   * @TODO remove this comment
-   */
   request = (async (...args: Parameters<LockManager['request']>) => {
     const [
       name,
       options = {
+        ifAvailable: false,
         mode: 'exclusive',
+        steal: false,
       },
       cb,
     ] = args;
 
-    const { mode } = options;
+    const { ifAvailable, mode, steal } = options;
+
+    const lock = {
+      name,
+      mode,
+      abortController: new AbortController(),
+    } as MockLock;
+
+    const heldLocksWithSameName = getOrCreateMapValue(
+      this.locks.held,
+      name,
+      new Map<LockMode, Set<MockLock>>(),
+    );
+
+    const exclusiveLocksWithSameName = getOrCreateMapValue(
+      heldLocksWithSameName,
+      'exclusive',
+      new Set<MockLock>(),
+    );
 
-    const lock = { name, mode, cb } as Lock;
+    const sharedLocksWithSameName = getOrCreateMapValue(
+      heldLocksWithSameName,
+      'shared',
+      new Set<MockLock>(),
+    );
 
     if (mode === 'shared') {
-      this.locks.held.add(lock);
+      sharedLocksWithSameName.add(lock);
 
-      await cb(lock);
+      try {
+        await cb(lock);
+      } finally {
+        sharedLocksWithSameName.delete(lock);
+      }
+
+      return;
+    }
+
+    // exclusive lock
+
+    // no locks with this name -> grant an exclusive lock
+    if (
+      exclusiveLocksWithSameName.size === 0 &&
+      sharedLocksWithSameName.size === 0
+    ) {
+      exclusiveLocksWithSameName.add(lock);
+
+      try {
+        await cb(lock);
+      } finally {
+        exclusiveLocksWithSameName.delete(lock);
+      }
 
       return;
     }
+
+    // steal any currently held locks
+    if (steal) {
+      for (const lock of sharedLocksWithSameName) {
+        lock.abortController.abort();
+      }
+
+      for (const lock of exclusiveLocksWithSameName) {
+        lock.abortController.abort();
+      }
+
+      sharedLocksWithSameName.clear();
+      exclusiveLocksWithSameName.clear();
+
+      exclusiveLocksWithSameName.add(lock);
+
+      try {
+        await cb(lock);
+      } finally {
+        exclusiveLocksWithSameName.delete(lock);
+      }
+
+      return;
+    }
+
+    // run the callback without waiting for the lock to be released
+    if (ifAvailable) {
+      // just run the callback without waiting for it
+      cb(null);
+
+      return;
+    }
+
+    // @TODO add the lock to the list of pending locks?
+
+    // it's an exclusive lock, so there's only one value
+    const currentLock = exclusiveLocksWithSameName.values().next()
+      .value as MockLock;
+
+    const { abortController: currentLockAbortController } = currentLock;
+
+    // wait for the current lock to be released
+    await new Promise<void>((resolve, reject) => {
+      currentLockAbortController.signal.onabort = () => resolve();
+
+      const { abortController: pendingLockAbortController } = lock;
+
+      // this allows the locking mechanism to release this lock
+      pendingLockAbortController.signal.onabort = () => reject();
+    });
+
+    // clear the exclusive locks
+    exclusiveLocksWithSameName.clear();
+
+    // grant our lock
+    exclusiveLocksWithSameName.add(lock);
+
+    try {
+      // run the callback
+      await cb(lock);
+    } finally {
+      exclusiveLocksWithSameName.delete(lock);
+    }
+
+    return;
   }) as LockManager['request'];
 }
 
diff --git a/tests/utils.ts b/tests/utils.ts
index e4874e4ef8..68f0867649 100644
--- a/tests/utils.ts
+++ b/tests/utils.ts
@@ -146,3 +146,25 @@ export function emitEvent(type: FiddleEvent, ...args: any[]) {
     },
   );
 }
+
+export function getOrCreateMapValue<T extends Map<unknown, unknown>>(
+  map: T,
+  key: MapKey<T>,
+  fallbackValue: MapValue<T>,
+): MapValue<T> {
+  if (!map.has(key)) {
+    map.set(key, fallbackValue);
+
+    return fallbackValue;
+  }
+
+  return map.get(key) as MapValue<T>;
+}
+
+type MapKey<T extends Map<unknown, unknown>> = T extends Map<infer I, unknown>
+  ? I
+  : never;
+
+type MapValue<T extends Map<unknown, unknown>> = T extends Map<unknown, infer I>
+  ? I
+  : never;