feat(waitForElementToBeRemoved): support passing an element directly (#460)

Author: kentcdodds

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

@@ -45,27 +45,35 @@ 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 element(s) given to waitForElementToBeRemoved are already removed. 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 element(s) given to waitForElementToBeRemoved are already removed. 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."`,
+    `"The element(s) given to waitForElementToBeRemoved are already removed. 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(
-    `"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."`,
+    `"The element(s) given to waitForElementToBeRemoved are already removed. waitForElementToBeRemoved requires that the element(s) exist(s) before waiting for removal."`,
   )
 })
 
@@ -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})
+})

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

@@ -2,23 +2,31 @@ import {waitFor} from './wait-for'
 
 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.
+// As the name implies, waitForElementToBeRemoved should check `present` --> `removed`
+function initialCheck(elements) {
+  if (isRemoved(elements)) {
+    throw new Error(
+      'The element(s) given to waitForElementToBeRemoved are already removed. waitForElementToBeRemoved requires that the element(s) exist(s) before waiting for removal.',
     )
   }
+}
 
-  // Check if the element is not present synchronously,
-  // As the name implies, waitForElementToBeRemoved should check `present` --> `removed`
-  if (isRemoved(callback())) {
-    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 waitFor(() => {
     let result
     try {