From 2b8e9e8633dc5ecf93198a2ce08901ee636e0f65 Mon Sep 17 00:00:00 2001
From: Patrick <patomation@gmail.com>
Date: Sun, 22 Sep 2024 12:32:56 -0700
Subject: [PATCH] fix: useKeyMap unmount removes keys (3778)

---
 .../src/editor/__tests__/hooks.spec.ts        | 31 +++++++++++++++++++
 packages/graphiql-react/src/editor/hooks.ts   | 10 ++++--
 2 files changed, 38 insertions(+), 3 deletions(-)
 create mode 100644 packages/graphiql-react/src/editor/__tests__/hooks.spec.ts

diff --git a/packages/graphiql-react/src/editor/__tests__/hooks.spec.ts b/packages/graphiql-react/src/editor/__tests__/hooks.spec.ts
new file mode 100644
index 00000000000..b7405efd467
--- /dev/null
+++ b/packages/graphiql-react/src/editor/__tests__/hooks.spec.ts
@@ -0,0 +1,31 @@
+import { renderHook } from '@testing-library/react';
+import { useKeyMap } from '../hooks';
+import { CodeMirrorEditor, KeyMap } from '../types';
+
+describe('hooks', () => {
+  describe('useKeyMap', () => {
+    it('works correctly', () => {
+      type KeyMap = Record<string, typeof jest.fn>;
+      let keys: KeyMap = {};
+      const editor: Pick<CodeMirrorEditor, 'addKeyMap' | 'removeKeyMap'> = {
+        addKeyMap: jest.fn(keyMap => {
+          keys = { ...keys, ...(keyMap as KeyMap) };
+        }),
+        removeKeyMap: jest.fn((key: string) => {
+          delete keys[key];
+        }),
+      };
+      const callback = jest.fn();
+      const { unmount } = renderHook(() =>
+        useKeyMap(
+          editor as unknown as CodeMirrorEditor,
+          ['foo', 'bar'],
+          callback,
+        ),
+      );
+      expect(Object.keys(keys).length).toBe(2);
+      unmount();
+      expect(Object.keys(keys).length).toBe(0);
+    });
+  });
+});
diff --git a/packages/graphiql-react/src/editor/hooks.ts b/packages/graphiql-react/src/editor/hooks.ts
index e946c59926b..8f56954f8b7 100644
--- a/packages/graphiql-react/src/editor/hooks.ts
+++ b/packages/graphiql-react/src/editor/hooks.ts
@@ -136,9 +136,12 @@ export function useKeyMap(
     if (!editor) {
       return;
     }
-    for (const key of keys) {
-      editor.removeKeyMap(key);
-    }
+    const handleRemoveKeys = () => {
+      for (const key of keys) {
+        editor.removeKeyMap(key);
+      }
+    };
+    handleRemoveKeys();
 
     if (callback) {
       const keyMap: Record<string, EmptyCallback> = {};
@@ -147,6 +150,7 @@ export function useKeyMap(
       }
       editor.addKeyMap(keyMap);
     }
+    return handleRemoveKeys;
   }, [editor, keys, callback]);
 }