fix(waitForElementToBeRemoved): Make initial check work with getBy*.

This commit includes a regression test. Fixes testing-library#1093.

Author: make-github-pseudonymous-again

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

@@ -80,6 +80,17 @@ test('requires an unempty array of elements to exist first (function form)', ()
   )
 })
 
+test('requires a getBy* query to not throw first', () => {
+  const {getByTestId} = renderIntoDocument(`
+  <div data-testid="div"></div>
+`)
+  return expect(
+    waitForElementToBeRemoved(() => getByTestId('non-existing-testid')),
+  ).rejects.toThrowErrorMatchingInlineSnapshot(
+    `The element(s) given to waitForElementToBeRemoved are already removed. waitForElementToBeRemoved requires that the element(s) exist(s) before waiting for removal.`,
+  )
+})
+
 test('after successful removal, fullfills promise with empty value (undefined)', () => {
   const {getByTestId} = renderIntoDocument(`
   <div data-testid="div"></div>

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

@@ -12,13 +12,28 @@ function initialCheck(elements) {
   }
 }
 
+function wrapFunctionCallback(callback) {
+  return () => {
+    try {
+      return callback()
+    } catch (error) {
+      if (error.name === 'TestingLibraryElementError') {
+        return null
+      }
+      throw error
+    }
+  }
+}
+
 async function waitForElementToBeRemoved(callback, options) {
   // created here so we get a nice stacktrace
   const timeoutError = new Error('Timed out in waitForElementToBeRemoved.')
-  if (typeof callback !== 'function') {
-    initialCheck(callback)
+  if (typeof callback === 'function') {
+    callback = wrapFunctionCallback(callback)
+  } else {
     const elements = Array.isArray(callback) ? callback : [callback]
     const getRemainingElements = elements.map(element => {
+      if (!element) return () => null
       let parent = element.parentElement
       if (parent === null) return () => null
       while (parent.parentElement) parent = parent.parentElement
@@ -30,15 +45,7 @@ async function waitForElementToBeRemoved(callback, options) {
   initialCheck(callback())
 
   return waitFor(() => {
-    let result
-    try {
-      result = callback()
-    } catch (error) {
-      if (error.name === 'TestingLibraryElementError') {
-        return undefined
-      }
-      throw error
-    }
+    const result = callback()
     if (!isRemoved(result)) {
       throw timeoutError
     }