diff --git a/src/Magento/FunctionalTestingFramework/Module/MagentoPwaWebDriver.php b/src/Magento/FunctionalTestingFramework/Module/MagentoPwaWebDriver.php
index 272f4206d..f2f818bd4 100644
--- a/src/Magento/FunctionalTestingFramework/Module/MagentoPwaWebDriver.php
+++ b/src/Magento/FunctionalTestingFramework/Module/MagentoPwaWebDriver.php
@@ -63,12 +63,13 @@ public function waitForPwaElementNotVisible($selector, $timeout = null)
 
         // Determine what type of Selector is used.
         // Then use the correct JavaScript to locate the Element.
-        if (\Codeception\Util\Locator::isXPath($selector)) {
+        if (\Codeception\Util\Locator::isCss($selector)) {
             $this->waitForLoadingMaskToDisappear($timeout);
-            $this->waitForJS("return !document.evaluate(`$selector`, document);", $timeout);
+            $this->waitForJS("return !document.querySelector(`$selector`);", $timeout);
         } else {
             $this->waitForLoadingMaskToDisappear($timeout);
-            $this->waitForJS("return !document.querySelector(`$selector`);", $timeout);
+            $this->waitForJS("return !document.evaluate(`$selector`, document, null, 
+            XPathResult.FIRST_ORDERED_NODE_TYPE, null).singleNodeValue;", $timeout);
         }
     }
 
@@ -84,15 +85,16 @@ public function waitForPwaElementNotVisible($selector, $timeout = null)
     public function waitForPwaElementVisible($selector, $timeout = null)
     {
         $timeout = $timeout ?? $this->_getConfig()['pageload_timeout'];
-        
+
         // Determine what type of Selector is used.
         // Then use the correct JavaScript to locate the Element.
-        if (\Codeception\Util\Locator::isXPath($selector)) {
+        if (\Codeception\Util\Locator::isCss($selector)) {
             $this->waitForLoadingMaskToDisappear($timeout);
-            $this->waitForJS("return !!document && !!document.evaluate(`$selector`, document);", $timeout);
+            $this->waitForJS("return !!document && !!document.querySelector(`$selector`);", $timeout);
         } else {
             $this->waitForLoadingMaskToDisappear($timeout);
-            $this->waitForJS("return !!document && !!document.querySelector(`$selector`);", $timeout);
+            $this->waitForJS("return !!document && !!document.evaluate(`$selector`, document, null, 
+            XPathResult.FIRST_ORDERED_NODE_TYPE, null).singleNodeValue;", $timeout);
         }
     }
 }