From 2dbc38c17a8389cb682cb6af588c066f03d6f686 Mon Sep 17 00:00:00 2001
From: Robin Malfait <malfait.robin@gmail.com>
Date: Thu, 17 Mar 2022 23:53:53 +0100
Subject: [PATCH] add tests to prove guarding against infinite loops is
 important (#1253)

---
 .../components/focus-trap/focus-trap.test.tsx | 79 ++++++++++++++++++-
 .../components/focus-trap/focus-trap.test.ts  | 68 ++++++++++++++++
 2 files changed, 143 insertions(+), 4 deletions(-)

diff --git a/packages/@headlessui-react/src/components/focus-trap/focus-trap.test.tsx b/packages/@headlessui-react/src/components/focus-trap/focus-trap.test.tsx
index 58c2c523b0..f342d1f7be 100644
--- a/packages/@headlessui-react/src/components/focus-trap/focus-trap.test.tsx
+++ b/packages/@headlessui-react/src/components/focus-trap/focus-trap.test.tsx
@@ -1,5 +1,5 @@
-import React, { useState, useRef } from 'react'
-import { render } from '@testing-library/react'
+import React, { useState, useRef, FocusEvent } from 'react'
+import { render, screen } from '@testing-library/react'
 
 import { FocusTrap } from './focus-trap'
 import { assertActiveElement } from '../../test-utils/accessibility-assertions'
@@ -7,13 +7,13 @@ import { suppressConsoleLogs } from '../../test-utils/suppress-console-logs'
 import { click, press, shift, Keys } from '../../test-utils/interactions'
 
 it('should focus the first focusable element inside the FocusTrap', () => {
-  let { getByText } = render(
+  render(
     <FocusTrap>
       <button>Trigger</button>
     </FocusTrap>
   )
 
-  assertActiveElement(getByText('Trigger'))
+  assertActiveElement(screen.getByText('Trigger'))
 })
 
 it('should focus the autoFocus element inside the FocusTrap if that exists', () => {
@@ -74,6 +74,7 @@ it('should warn when there is no focusable element inside the FocusTrap', () =>
   }
   render(<Example />)
   expect(spy.mock.calls[0][0]).toBe('There are no focusable elements inside the <FocusTrap />')
+  spy.mockReset()
 })
 
 it(
@@ -325,3 +326,73 @@ it('should be possible skip disabled elements within the focus trap', async () =
   await press(Keys.Tab)
   assertActiveElement(document.getElementById('item-a'))
 })
+
+it('should try to focus all focusable items (and fail)', async () => {
+  let spy = jest.spyOn(console, 'warn').mockImplementation(jest.fn())
+  let focusHandler = jest.fn()
+  function handleFocus(e: FocusEvent) {
+    let target = e.target as HTMLElement
+    focusHandler(target.id)
+    screen.getByText('After')?.focus()
+  }
+
+  render(
+    <>
+      <button id="before">Before</button>
+      <FocusTrap>
+        <button id="item-a" onFocus={handleFocus}>
+          Item A
+        </button>
+        <button id="item-b" onFocus={handleFocus}>
+          Item B
+        </button>
+        <button id="item-c" onFocus={handleFocus}>
+          Item C
+        </button>
+        <button id="item-d" onFocus={handleFocus}>
+          Item D
+        </button>
+      </FocusTrap>
+      <button>After</button>
+    </>
+  )
+
+  expect(focusHandler.mock.calls).toEqual([['item-a'], ['item-b'], ['item-c'], ['item-d']])
+  expect(spy).toHaveBeenCalledWith('There are no focusable elements inside the <FocusTrap />')
+  spy.mockReset()
+})
+
+it('should end up at the last focusable element', async () => {
+  let spy = jest.spyOn(console, 'warn').mockImplementation(jest.fn())
+
+  let focusHandler = jest.fn()
+  function handleFocus(e: FocusEvent) {
+    let target = e.target as HTMLElement
+    focusHandler(target.id)
+    screen.getByText('After')?.focus()
+  }
+
+  render(
+    <>
+      <button id="before">Before</button>
+      <FocusTrap>
+        <button id="item-a" onFocus={handleFocus}>
+          Item A
+        </button>
+        <button id="item-b" onFocus={handleFocus}>
+          Item B
+        </button>
+        <button id="item-c" onFocus={handleFocus}>
+          Item C
+        </button>
+        <button id="item-d">Item D</button>
+      </FocusTrap>
+      <button>After</button>
+    </>
+  )
+
+  expect(focusHandler.mock.calls).toEqual([['item-a'], ['item-b'], ['item-c']])
+  assertActiveElement(screen.getByText('Item D'))
+  expect(spy).not.toHaveBeenCalled()
+  spy.mockReset()
+})
diff --git a/packages/@headlessui-vue/src/components/focus-trap/focus-trap.test.ts b/packages/@headlessui-vue/src/components/focus-trap/focus-trap.test.ts
index 8d6707250f..1b24930ce7 100644
--- a/packages/@headlessui-vue/src/components/focus-trap/focus-trap.test.ts
+++ b/packages/@headlessui-vue/src/components/focus-trap/focus-trap.test.ts
@@ -124,6 +124,7 @@ it('should warn when there is no focusable element inside the FocusTrap', async
   await new Promise<void>(nextTick)
 
   expect(spy.mock.calls[0][0]).toBe('There are no focusable elements inside the <FocusTrap />')
+  spy.mockReset()
 })
 
 it(
@@ -379,3 +380,70 @@ it('should be possible skip disabled elements within the focus trap', async () =
   await press(Keys.Tab)
   assertActiveElement(document.getElementById('item-a'))
 })
+
+it('should try to focus all focusable items in order (and fail)', async () => {
+  let spy = jest.spyOn(console, 'warn').mockImplementation(jest.fn())
+  let focusHandler = jest.fn()
+
+  renderTemplate({
+    template: html`
+      <div>
+        <button id="before">Before</button>
+        <FocusTrap>
+          <button id="item-a" @focus="handleFocus">Item A</button>
+          <button id="item-b" @focus="handleFocus">Item B</button>
+          <button id="item-c" @focus="handleFocus">Item C</button>
+          <button id="item-d" @focus="handleFocus">Item D</button>
+        </FocusTrap>
+        <button>After</button>
+      </div>
+    `,
+    setup() {
+      return {
+        handleFocus(e: Event) {
+          let target = e.target as HTMLElement
+          focusHandler(target.id)
+          getByText('After')?.focus()
+        },
+      }
+    },
+  })
+
+  expect(focusHandler.mock.calls).toEqual([['item-a'], ['item-b'], ['item-c'], ['item-d']])
+  expect(spy).toHaveBeenCalledWith('There are no focusable elements inside the <FocusTrap />')
+  spy.mockReset()
+})
+
+it('should end up at the last focusable element', async () => {
+  let spy = jest.spyOn(console, 'warn').mockImplementation(jest.fn())
+  let focusHandler = jest.fn()
+
+  renderTemplate({
+    template: html`
+      <div>
+        <button id="before">Before</button>
+        <FocusTrap>
+          <button id="item-a" @focus="handleFocus">Item A</button>
+          <button id="item-b" @focus="handleFocus">Item B</button>
+          <button id="item-c" @focus="handleFocus">Item C</button>
+          <button id="item-d">Item D</button>
+        </FocusTrap>
+        <button>After</button>
+      </div>
+    `,
+    setup() {
+      return {
+        handleFocus(e: Event) {
+          let target = e.target as HTMLElement
+          focusHandler(target.id)
+          getByText('After')?.focus()
+        },
+      }
+    },
+  })
+
+  expect(focusHandler.mock.calls).toEqual([['item-a'], ['item-b'], ['item-c']])
+  assertActiveElement(getByText('Item D'))
+  expect(spy).not.toHaveBeenCalled()
+  spy.mockReset()
+})