From a8061e8638b20a6c72fb84977cb5a17d76adf4a0 Mon Sep 17 00:00:00 2001
From: "Eric Wong (PAX)" <eriwong@microsoft.com>
Date: Mon, 14 Sep 2020 17:16:34 -0700
Subject: [PATCH] Fix Race condition in tests involving context menu

---
 .../CalculatorResults.cs                      |  4 +-
 .../WindowsDriverExtensions.cs                | 43 +++++++++++++++++++
 2 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/src/CalculatorUITestFramework/CalculatorResults.cs b/src/CalculatorUITestFramework/CalculatorResults.cs
index 9371fa11..6e6815d9 100644
--- a/src/CalculatorUITestFramework/CalculatorResults.cs
+++ b/src/CalculatorUITestFramework/CalculatorResults.cs
@@ -13,8 +13,8 @@ namespace CalculatorUITestFramework
         private WindowsElement CalculatorAlwaysOnTopResults => this.session.TryFindElementByAccessibilityId("CalculatorAlwaysOnTopResults");
         private WindowsElement CalculatorResult => this.session.TryFindElementByAccessibilityId("CalculatorResults");
         private WindowsElement CalculatorExpression => this.session.TryFindElementByAccessibilityId("CalculatorExpression");
-        private WindowsElement MenuItemCopy => this.session.TryFindElementByAccessibilityId("CopyMenuItem");
-        private WindowsElement MenuItemPaste => this.session.TryFindElementByAccessibilityId("PasteMenuItem");
+        private WindowsElement MenuItemCopy => this.session.WaitForElementByAccessibilityId("CopyMenuItem");
+        private WindowsElement MenuItemPaste => this.session.WaitForElementByAccessibilityId("PasteMenuItem");
 
         /// <summary>
         /// Gets the text from the display control in AoT mode and removes the narrator text that is not displayed in the UI.
diff --git a/src/CalculatorUITestFramework/WindowsDriverExtensions.cs b/src/CalculatorUITestFramework/WindowsDriverExtensions.cs
index 399584a5..98bc8c0d 100644
--- a/src/CalculatorUITestFramework/WindowsDriverExtensions.cs
+++ b/src/CalculatorUITestFramework/WindowsDriverExtensions.cs
@@ -1,8 +1,11 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+using Microsoft.VisualStudio.TestTools.UnitTesting.Logging;
 using OpenQA.Selenium;
 using OpenQA.Selenium.Appium.Windows;
+using System.Diagnostics;
+using System.Threading;
 
 namespace CalculatorUITestFramework
 {
@@ -68,5 +71,45 @@ namespace CalculatorUITestFramework
             // process id (any entry of allWindowHandles)
             driver.SwitchTo().Window(allWindowHandles[0]);
         }
+
+        /// <summary>
+        /// Waits for an element to be created.
+        /// </summary>
+        /// <param name="driver">this</param>
+        /// <param name="id">the automation id</param>
+        /// <param name="timeout">optional timeout in ms</param>
+        /// <returns>the element with the matching automation id</returns>
+        public static WindowsElement WaitForElementByAccessibilityId(this WindowsDriver<WindowsElement> driver, string id, int timeout = 1000)
+        {
+            Stopwatch timer = new Stopwatch();
+            timer.Reset();
+            timer.Start();
+            while (timer.ElapsedMilliseconds < timeout)
+            {
+                try
+                {
+                    var element  = driver.TryFindElementByAccessibilityId(id);
+                    return element;
+                }
+                catch(WebDriverException ex)
+                {
+                    if (ex.Message.Contains("An element could not be located on the page using the given search parameters"))
+                    {
+                        Logger.LogMessage("Element not found. Waiting for 10ms in WaitForElementByAccessibilityId");
+                    }
+                    else
+                    {
+                        throw;
+                    }
+                }
+
+                Logger.LogMessage("Waiting for 10ms in WaitForElementByAccessibilityId");
+                Thread.Sleep(10);
+            }
+            timer.Stop();
+
+            // one last attempt. Throws the not found exception if this fails
+            return driver.TryFindElementByAccessibilityId(id);
+        }
     }
 }