Cycle focus between the body and page tab sequence

See testing-library#365 for context

src/__tests__/tab.js

@@ -111,7 +111,11 @@ test('should cycle elements in document tab order', () => {
 
   userEvent.tab()
 
-  // cycle goes back to first element
+  // cycle goes back to the body
+  expect(document.body).toHaveFocus()
+
+  userEvent.tab()
+
   expect(checkbox).toHaveFocus()
 })
 
@@ -127,14 +131,27 @@ test('should go backwards when shift = true', () => {
     '[data-testid="element"]',
   )
 
-  radio.focus()
+  expect(document.body).toHaveFocus()
+
+  userEvent.tab({shift: true})
+
+  expect(number).toHaveFocus()
+
+  userEvent.tab({shift: true})
+
+  expect(radio).toHaveFocus()
 
   userEvent.tab({shift: true})
 
   expect(checkbox).toHaveFocus()
 
   userEvent.tab({shift: true})
 
+  // cycle goes back to the body
+  expect(document.body).toHaveFocus()
+
+  userEvent.tab({shift: true})
+
   expect(number).toHaveFocus()
 })
 
@@ -164,6 +181,10 @@ test('should respect tabindex, regardless of dom position', () => {
 
   userEvent.tab()
 
+  expect(document.body).toHaveFocus()
+
+  userEvent.tab()
+
   expect(radio).toHaveFocus()
 })
 
@@ -193,6 +214,10 @@ test('should respect tab index order, then DOM order', () => {
 
   userEvent.tab()
 
+  expect(document.body).toHaveFocus()
+
+  userEvent.tab()
+
   expect(checkbox).toHaveFocus()
 })
 

src/tab.js

@@ -3,6 +3,22 @@ import {getActiveElement, FOCUSABLE_SELECTOR} from './utils'
 import {focus} from './focus'
 import {blur} from './blur'
 
+function getNextElement(currentIndex, shift, elements, focusTrap) {
+  if (focusTrap === document && currentIndex === 0 && shift) {
+    return document.body
+  } else if (
+    focusTrap === document &&
+    currentIndex === elements.length - 1 &&
+    !shift
+  ) {
+    return document.body
+  } else {
+    const nextIndex = shift ? currentIndex - 1 : currentIndex + 1
+    const defaultIndex = shift ? elements.length - 1 : 0
+    return elements[nextIndex] || elements[defaultIndex]
+  }
+}
+
 function tab({shift = false, focusTrap} = {}) {
   const previousElement = getActiveElement(focusTrap?.ownerDocument ?? document)
 
@@ -57,10 +73,7 @@ function tab({shift = false, focusTrap} = {}) {
     el => el === el.ownerDocument.activeElement,
   )
 
-  const nextIndex = shift ? index - 1 : index + 1
-  const defaultIndex = shift ? prunedElements.length - 1 : 0
-
-  const nextElement = prunedElements[nextIndex] || prunedElements[defaultIndex]
+  const nextElement = getNextElement(index, shift, prunedElements, focusTrap)
 
   const shiftKeyInit = {
     key: 'Shift',