From 8dbcbb354f974009ef42d9795ba789dfd8919f01 Mon Sep 17 00:00:00 2001
From: Priyansh Garg <priyanshgarg30@gmail.com>
Date: Thu, 14 Dec 2023 15:56:45 +0530
Subject: [PATCH] Fix waitFor commands not retrying for stale elements. (#3978)

---
 lib/api/element-commands/_waitForDisplayed.js |  6 ++++-
 .../element/testWaitForElementVisible.js      | 27 +++++++++----------
 2 files changed, 18 insertions(+), 15 deletions(-)

diff --git a/lib/api/element-commands/_waitForDisplayed.js b/lib/api/element-commands/_waitForDisplayed.js
index d2b91128f4..3fc575571d 100644
--- a/lib/api/element-commands/_waitForDisplayed.js
+++ b/lib/api/element-commands/_waitForDisplayed.js
@@ -10,7 +10,11 @@ class WaitForDisplayed extends WaitForElement {
   }
 
   setupActions() {
-    const isResultStale = (result) => this.transport.staleElementReference(result);
+    const isResultStale = (result) => {
+      const errorResponse = this.transport.getErrorResponse(result);
+
+      return this.transport.staleElementReference(errorResponse);
+    };
     const validate = (result) => this.isResultSuccess(result);
     const successHandler = (result) => this.protocolActionHandler(result);
 
diff --git a/test/src/api/commands/element/testWaitForElementVisible.js b/test/src/api/commands/element/testWaitForElementVisible.js
index e5966c16c3..b14948c7e7 100644
--- a/test/src/api/commands/element/testWaitForElementVisible.js
+++ b/test/src/api/commands/element/testWaitForElementVisible.js
@@ -160,41 +160,40 @@ describe('waitForElementVisible', function () {
         postdata: '{"using":"css selector","value":"#stale-element"}',
         method: 'POST',
         response: JSON.stringify({
-          status: 0,
-          state: 'success',
           value: [{'element-6066-11e4-a52e-4f735466cecf': '99'}]
         })
-      })
+      }, true)
       .addMock({
         url: '/wd/hub/session/1352110219202/execute/sync',
         method: 'POST',
         response: JSON.stringify({
-          sessionId: '1352110219202',
-          state: 'stale element reference',
-          status: 10
+          value: {
+            error: 'stale element reference'
+          }
+        })
+      }, true)
+      .addMock({
+        url: '/wd/hub/session/1352110219202/elements',
+        postdata: '{"using":"css selector","value":"#stale-element"}',
+        method: 'POST',
+        response: JSON.stringify({
+          value: [{'element-6066-11e4-a52e-4f735466cecf': '88'}]
         })
       }, true)
       .addMock({
         url: '/wd/hub/session/1352110219202/execute/sync',
         method: 'POST',
         response: JSON.stringify({
-          state: 'success',
-          status: 0,
           value: true
         })
       }, true);
 
     this.client.api.waitForElementVisible('#stale-element', 110, 10, function callback(result, instance) {
-      assert.strictEqual(instance.elementId, '99');
+      assert.strictEqual(instance.elementId, '88');  // element id refreshed
       assert.strictEqual(result.value, true);
     });
 
     return this.client.start(function (err) {
-      MockServer.removeMock({
-        url: '/wd/hub/session/1352110219202/elements',
-        method: 'POST'
-      });
-
       if (err) {
         throw err;
       }