feat(waitForElementToBeRemoved): support passing an element directly

Author: kentcdodds

src/__tests__/wait-for-element-to-be-removed.js

@@ -45,23 +45,31 @@ test('resolves on mutation if callback throws an error', async () => {
   await waitForElementToBeRemoved(() => getByTestId('div'), {timeout: 100})
 })
 
-test('requires a function as the first parameter', () => {
+test('requires an element to exist first', () => {
   return expect(
-    waitForElementToBeRemoved(),
+    waitForElementToBeRemoved(null),
   ).rejects.toThrowErrorMatchingInlineSnapshot(
-    `"waitForElementToBeRemoved requires a callback as the first parameter"`,
+    `"The callback function which was passed did not return an element or non-empty array of elements. waitForElementToBeRemoved requires that the element(s) exist(s) before waiting for removal."`,
   )
 })
 
-test('requires an element to exist first', () => {
+test('requires an unempty array of elements to exist first', () => {
+  return expect(
+    waitForElementToBeRemoved([]),
+  ).rejects.toThrowErrorMatchingInlineSnapshot(
+    `"The callback function which was passed did not return an element or non-empty array of elements. waitForElementToBeRemoved requires that the element(s) exist(s) before waiting for removal."`,
+  )
+})
+
+test('requires an element to exist first (function form)', () => {
   return expect(
     waitForElementToBeRemoved(() => null),
   ).rejects.toThrowErrorMatchingInlineSnapshot(
     `"The callback function which was passed did not return an element or non-empty array of elements. waitForElementToBeRemoved requires that the element(s) exist(s) before waiting for removal."`,
   )
 })
 
-test('requires an unempty array of elements to exist first', () => {
+test('requires an unempty array of elements to exist first (function form)', () => {
   return expect(
     waitForElementToBeRemoved(() => []),
   ).rejects.toThrowErrorMatchingInlineSnapshot(
@@ -132,3 +140,37 @@ test('rethrows non-testing-lib errors', () => {
     }),
   ).rejects.toBe(error)
 })
+
+test('accepts an element as an argument and waits for it to be removed from its top-most parent', async () => {
+  const {queryByTestId} = renderIntoDocument(`
+    <div data-testid="div"></div>
+  `)
+  const div = queryByTestId('div')
+  setTimeout(() => {
+    div.parentElement.removeChild(div)
+  }, 20)
+
+  await waitForElementToBeRemoved(div, {timeout: 200})
+})
+
+test('accepts an array of elements as an argument and waits for those elements to be removed from their top-most parent', async () => {
+  const {queryAllByTestId} = renderIntoDocument(`
+    <div>
+      <div>
+        <div data-testid="div"></div>
+      </div>
+      <div>
+        <div data-testid="div"></div>
+      </div>
+    </div>
+  `)
+  const [div1, div2] = queryAllByTestId('div')
+  setTimeout(() => {
+    div1.parentElement.removeChild(div1)
+  }, 20)
+
+  setTimeout(() => {
+    div2.parentElement.removeChild(div2)
+  }, 50)
+  await waitForElementToBeRemoved([div1, div2], {timeout: 200})
+})

src/wait-for-element-to-be-removed.js

@@ -2,22 +2,30 @@ import {wait} from './wait'
 
 const isRemoved = result => !result || (Array.isArray(result) && !result.length)
 
-async function waitForElementToBeRemoved(callback, options) {
-  if (!callback) {
-    return Promise.reject(
-      new Error(
-        'waitForElementToBeRemoved requires a callback as the first parameter',
-      ),
-    )
-  }
-
-  // Check if the element is not present synchronously,
-  // As the name implies, waitForElementToBeRemoved should check `present` --> `removed`
-  if (isRemoved(callback())) {
+// Check if the element is not present.
+// As the name implies, waitForElementToBeRemoved should check `present` --> `removed`
+function initialCheck(elements) {
+  if (isRemoved(elements)) {
     throw new Error(
       'The callback function which was passed did not return an element or non-empty array of elements. waitForElementToBeRemoved requires that the element(s) exist(s) before waiting for removal.',
     )
   }
+}
+
+async function waitForElementToBeRemoved(callback, options) {
+  if (typeof callback !== 'function') {
+    // await waitForElementToBeRemoved(getAllByText('Hello'))
+    initialCheck(callback)
+    const elements = Array.isArray(callback) ? callback : [callback]
+    const getRemainingElements = elements.map(element => {
+      let parent = element.parentElement
+      while (parent.parentElement) parent = parent.parentElement
+      return () => (parent.contains(element) ? element : null)
+    })
+    callback = () => getRemainingElements.map(c => c()).filter(Boolean)
+  }
+
+  initialCheck(callback())
 
   return wait(() => {
     let result