From d5b3f35775a5c80ddd7219405032ca6307a6d7e8 Mon Sep 17 00:00:00 2001
From: make-github-pseudonymous-again
 <5165674+make-github-pseudonymous-again@users.noreply.github.com>
Date: Fri, 21 Jan 2022 17:16:41 +0100
Subject: [PATCH] fix(waitForElementToBeRemoved): Make initial check work with
 getBy*.

This commit includes a regression test. Fixes #1093.
---
 .../wait-for-element-to-be-removed.js         | 11 +++++++
 src/wait-for-element-to-be-removed.js         | 29 ++++++++++++-------
 2 files changed, 29 insertions(+), 11 deletions(-)

diff --git a/src/__tests__/wait-for-element-to-be-removed.js b/src/__tests__/wait-for-element-to-be-removed.js
index 775b8b82..5e170c10 100644
--- a/src/__tests__/wait-for-element-to-be-removed.js
+++ b/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>
diff --git a/src/wait-for-element-to-be-removed.js b/src/wait-for-element-to-be-removed.js
index 2b5bcb9a..5a3e4629 100644
--- a/src/wait-for-element-to-be-removed.js
+++ b/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
     }