From fb31fd7548a7495b80d0217797d793376b76555e Mon Sep 17 00:00:00 2001
From: DudaGod <dudkevich@yandex-team.ru>
Date: Tue, 26 Dec 2023 17:50:10 +0300
Subject: [PATCH] fix: correctly switch between iframes on disable animation

---
 src/browser/existing-browser.js      |  6 ++---
 test/src/browser/existing-browser.js | 37 +++++++++++++++++++---------
 2 files changed, 28 insertions(+), 15 deletions(-)

diff --git a/src/browser/existing-browser.js b/src/browser/existing-browser.js
index 9e56e786b..f17909279 100644
--- a/src/browser/existing-browser.js
+++ b/src/browser/existing-browser.js
@@ -300,12 +300,12 @@ module.exports = class ExistingBrowser extends Browser {
             for (const iframe of iframes) {
                 await this._session.switchToFrame(iframe);
                 await cb();
-            }
-        } finally {
-            if (!_.isEmpty(iframes)) {
                 // switchToParentFrame does not work in ios - https://github.com/appium/appium/issues/14882
                 await this._session.switchToFrame(null);
             }
+        } catch (e) {
+            await this._session.switchToFrame(null);
+            throw e;
         }
     }
 
diff --git a/test/src/browser/existing-browser.js b/test/src/browser/existing-browser.js
index 62d8bb911..9ddfa3dd4 100644
--- a/test/src/browser/existing-browser.js
+++ b/test/src/browser/existing-browser.js
@@ -722,20 +722,33 @@ describe("ExistingBrowser", () => {
             );
         });
 
-        it("should disable animations if 'disableAnimation: true' and 'automationProtocol: webdriver'", async () => {
-            const clientBridge = stubClientBridge_();
-            const browser = await initBrowser_(mkBrowser_({ automationProtocol: "webdriver" }));
-            const iframeElement = { "element-12345": "67890_element_1" };
-            browser.publicAPI.findElements.withArgs("css selector", "iframe").resolves([iframeElement]);
+        describe("'disableAnimation: true' and 'automationProtocol: webdriver'", () => {
+            it("should disable animations", async () => {
+                const clientBridge = stubClientBridge_();
+                const browser = await initBrowser_(mkBrowser_({ automationProtocol: "webdriver" }));
+                const iframeElement1 = { "element-12345": "67890_element_1" };
+                const iframeElement2 = { "element-54321": "09876_element_2" };
+                browser.publicAPI.findElements
+                    .withArgs("css selector", "iframe")
+                    .resolves([iframeElement1, iframeElement2]);
 
-            await browser.prepareScreenshot(".selector", { disableAnimation: true });
+                await browser.prepareScreenshot(".selector", { disableAnimation: true });
 
-            assert.calledWith(clientBridge.call, "prepareScreenshot", [
-                ".selector",
-                sinon.match({ disableAnimation: true }),
-            ]);
-            assert.calledWith(browser.publicAPI.switchToFrame, iframeElement);
-            assert.calledWith(clientBridge.call, "disableFrameAnimations");
+                assert.callOrder(
+                    clientBridge.call.withArgs("prepareScreenshot", [
+                        ".selector",
+                        sinon.match({ disableAnimation: true }),
+                    ]),
+
+                    browser.publicAPI.switchToFrame.withArgs(iframeElement1),
+                    clientBridge.call.withArgs("disableFrameAnimations"),
+                    browser.publicAPI.switchToFrame.withArgs(null),
+
+                    browser.publicAPI.switchToFrame.withArgs(iframeElement2),
+                    clientBridge.call.withArgs("disableFrameAnimations"),
+                    browser.publicAPI.switchToFrame.withArgs(null),
+                );
+            });
         });
 
         it("should not disable iframe animations if 'disableAnimation: true' and 'automationProtocol: devtools'", async () => {