From cd5924c4904e831ea38226495dc607659f67c3f8 Mon Sep 17 00:00:00 2001
From: Robin Malfait <malfait.robin@gmail.com>
Date: Tue, 10 Jan 2023 14:41:23 +0100
Subject: [PATCH] add tests to guarantee `FocusTrap` with a single element
 works as expected

---
 .../components/focus-trap/focus-trap.test.tsx | 94 +++++++++++++++++++
 .../components/focus-trap/focus-trap.test.ts  | 54 +++++++++++
 2 files changed, 148 insertions(+)

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 7f43d5ccd9..054ed6d9b8 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
@@ -108,6 +108,50 @@ it('should warn when there is no focusable element inside the FocusTrap', async
   spy.mockReset()
 })
 
+it(
+  'should not be possible to programmatically escape the focus trap (if there is only 1 focusable element)',
+  suppressConsoleLogs(async () => {
+    function Example() {
+      return (
+        <>
+          <input id="a" autoFocus />
+
+          <FocusTrap>
+            <input id="b" />
+          </FocusTrap>
+        </>
+      )
+    }
+
+    render(<Example />)
+
+    await nextFrame()
+
+    let [a, b] = Array.from(document.querySelectorAll('input'))
+
+    // Ensure that input-b is the active element
+    assertActiveElement(b)
+
+    // Tab to the next item
+    await press(Keys.Tab)
+
+    // Ensure that input-b is still the active element
+    assertActiveElement(b)
+
+    // Try to move focus
+    a?.focus()
+
+    // Ensure that input-b is still the active element
+    assertActiveElement(b)
+
+    // Click on an element within the FocusTrap
+    await click(b)
+
+    // Ensure that input-b is the active element
+    assertActiveElement(b)
+  })
+)
+
 it(
   'should not be possible to programmatically escape the focus trap',
   suppressConsoleLogs(async () => {
@@ -214,6 +258,56 @@ it('should restore the previously focused element, before entering the FocusTrap
   assertActiveElement(document.getElementById('item-2'))
 })
 
+it('should stay in the FocusTrap when using `tab`, if there is only 1 focusable element', async () => {
+  render(
+    <>
+      <button>Before</button>
+      <FocusTrap>
+        <button id="item-a">Item A</button>
+      </FocusTrap>
+      <button>After</button>
+    </>
+  )
+
+  await nextFrame()
+
+  // Item A should be focused because the FocusTrap will focus the first item
+  assertActiveElement(document.getElementById('item-a'))
+
+  // Next
+  await press(Keys.Tab)
+  assertActiveElement(document.getElementById('item-a'))
+
+  // Next
+  await press(Keys.Tab)
+  assertActiveElement(document.getElementById('item-a'))
+})
+
+it('should stay in the FocusTrap when using `shift+tab`, if there is only 1 focusable element', async () => {
+  render(
+    <>
+      <button>Before</button>
+      <FocusTrap>
+        <button id="item-a">Item A</button>
+      </FocusTrap>
+      <button>After</button>
+    </>
+  )
+
+  await nextFrame()
+
+  // Item A should be focused because the FocusTrap will focus the first item
+  assertActiveElement(document.getElementById('item-a'))
+
+  // Previous (loop around!)
+  await press(shift(Keys.Tab))
+  assertActiveElement(document.getElementById('item-a'))
+
+  // Previous
+  await press(shift(Keys.Tab))
+  assertActiveElement(document.getElementById('item-a'))
+})
+
 it('should be possible tab to the next focusable element within the focus trap', async () => {
   render(
     <>
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 c7996c7ee3..f3c4c58cef 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
@@ -227,6 +227,60 @@ it('should restore the previously focused element, before entering the FocusTrap
   assertActiveElement(document.getElementById('item-2'))
 })
 
+it('should stay in the FocusTrap when using `tab`, if there is only 1 focusable element', async () => {
+  renderTemplate({
+    template: html`
+      <div>
+        <button>Before</button>
+        <FocusTrap>
+          <button id="item-a">Item A</button>
+        </FocusTrap>
+        <button>After</button>
+      </div>
+    `,
+  })
+
+  await nextFrame()
+
+  // Item A should be focused because the FocusTrap will focus the first item
+  assertActiveElement(document.getElementById('item-a'))
+
+  // Next
+  await press(Keys.Tab)
+  assertActiveElement(document.getElementById('item-a'))
+
+  // Next
+  await press(Keys.Tab)
+  assertActiveElement(document.getElementById('item-a'))
+})
+
+it('should stay in the FocusTrap when using `shift+tab`, if there is only 1 focusable element', async () => {
+  renderTemplate({
+    template: html`
+      <div>
+        <button>Before</button>
+        <FocusTrap>
+          <button id="item-a">Item A</button>
+        </FocusTrap>
+        <button>After</button>
+      </div>
+    `,
+  })
+
+  await nextFrame()
+
+  // Item A should be focused because the FocusTrap will focus the first item
+  assertActiveElement(document.getElementById('item-a'))
+
+  // Previous (loop around!)
+  await press(shift(Keys.Tab))
+  assertActiveElement(document.getElementById('item-a'))
+
+  // Previous
+  await press(shift(Keys.Tab))
+  assertActiveElement(document.getElementById('item-a'))
+})
+
 it('should be possible to tab to the next focusable element within the focus trap', async () => {
   renderTemplate(
     html`