Replace HostWaitHelper (integration tests) with asynchronous code

src/VisualStudio/IntegrationTest/TestUtilities/Helper.cs

@@ -3,6 +3,7 @@
 using System;
 using System.Runtime.InteropServices;
 using System.Threading;
+using System.Threading.Tasks;
 using UIAutomationClient;
 
 namespace Microsoft.VisualStudio.IntegrationTest.Utilities
@@ -74,6 +75,32 @@ public static T Retry<T>(Func<T> action, TimeSpan delay)
                 delay);
         }
 
+        /// <summary>
+        /// This method will retry the asynchronous action represented by <paramref name="action"/>,
+        /// waiting for <paramref name="delay"/> time after each retry. If a given retry returns a value 
+        /// other than the default value of <typeparamref name="T"/>, this value is returned.
+        /// </summary>
+        /// <param name="action">the asynchronous action to retry</param>
+        /// <param name="delay">the amount of time to wait between retries</param>
+        /// <typeparam name="T">type of return value</typeparam>
+        /// <returns>the return value of <paramref name="action"/></returns>
+        public static Task<T> RetryAsync<T>(Func<Task<T>> action, TimeSpan delay)
+        {
+            return RetryAsyncHelper(async () =>
+            {
+                try
+                {
+                    return await action();
+                }
+                catch (COMException)
+                {
+                    // Devenv can throw COMExceptions if it's busy when we make DTE calls.
+                    return default;
+                }
+            },
+                delay);
+        }
+
         /// <summary>
         /// This method will retry the action represented by the 'action' argument,
         /// milliseconds, waiting 'delay' milliseconds after each retry and will swallow all exceptions. 
@@ -113,5 +140,19 @@ private static T RetryHelper<T>(Func<T> action, TimeSpan delay)
                 System.Threading.Thread.Sleep(delay);
             }
         }
+
+        private static async Task<T> RetryAsyncHelper<T>(Func<Task<T>> action, TimeSpan delay)
+        {
+            while (true)
+            {
+                var retval = await action().ConfigureAwait(true);
+                if (!Equals(default(T), retval))
+                {
+                    return retval;
+                }
+
+                await Task.Delay(delay).ConfigureAwait(true);
+            }
+        }
     }
 }

src/VisualStudio/IntegrationTest/TestUtilities/HostWaitHelper.cs

-// Copyright (c) Microsoft.  All Rights Reserved.  Licensed under the Apache License, Version 2.0.  See License.txt in the project root for license information.
-
-using System;
-using System.Collections.Generic;
-using System.Linq;
-using System.Threading.Tasks;
-using System.Windows;
-using System.Windows.Threading;
-
-namespace Microsoft.VisualStudio.IntegrationTest.Utilities
-{
-    public static class HostWaitHelper
-    {
-        public static void WaitForDispatchedOperationsToComplete(DispatcherPriority priority)
-            => Dispatcher.CurrentDispatcher.Invoke(() => { }, priority);
-
-        public static void PumpingWait(Task task)
-        {
-            var tasks = new[] { task };
-            PumpingWaitAll(tasks);
-        }
-
-        public static T PumpingWaitResult<T>(Task<T> task)
-        {
-            PumpingWait(task);
-            return task.Result;
-        }
-
-        public static void PumpingWaitAll(IEnumerable<Task> tasks)
-        {
-            var smallTimeout = TimeSpan.FromMilliseconds(10);
-            var taskArray = tasks.ToArray();
-            var done = false;
-            while (!done)
-            {
-                done = Task.WaitAll(taskArray, smallTimeout);
-                if (!done)
-                {
-                    WaitForDispatchedOperationsToComplete(DispatcherPriority.ApplicationIdle);
-                }
-            }
-
-            foreach (var task in tasks)
-            {
-                if (task.Exception != null)
-                {
-                    throw task.Exception;
-                }
-            }
-        }
-    }
-}

src/VisualStudio/IntegrationTest/TestUtilities/InProcess/Editor_InProc.cs

@@ -2,32 +2,31 @@
 
 using System;
 using System.Collections.Generic;
-using System.Collections.Immutable;
 using System.ComponentModel;
 using System.ComponentModel.Design;
 using System.Linq;
 using System.Runtime.InteropServices;
 using System.Threading;
+using System.Threading.Tasks;
 using System.Windows;
 using System.Windows.Controls;
 using System.Windows.Forms;
-using Microsoft.CodeAnalysis.Editor.Implementation.Highlighting;
 using System.Windows.Media;
+using Microsoft.CodeAnalysis.Editor.Implementation.Highlighting;
 using Microsoft.CodeAnalysis.Editor.Shared.Extensions;
-using Microsoft.CodeAnalysis.Text;
 using Microsoft.VisualStudio.IntegrationTest.Utilities.Common;
 using Microsoft.VisualStudio.Language.Intellisense;
 using Microsoft.VisualStudio.Shell.Interop;
 using Microsoft.VisualStudio.Text;
+using Microsoft.VisualStudio.Text.Classification;
 using Microsoft.VisualStudio.Text.Editor;
-using Microsoft.VisualStudio.Text.Operations;
 using Microsoft.VisualStudio.Text.Outlining;
 using Microsoft.VisualStudio.Text.Tagging;
 using Microsoft.VisualStudio.TextManager.Interop;
-using Microsoft.VisualStudio.Text.Classification;
 using UIAutomationClient;
 using AutomationElementIdentifiers = System.Windows.Automation.AutomationElementIdentifiers;
 using ControlType = System.Windows.Automation.ControlType;
+using ThreadHelper = Microsoft.VisualStudio.Shell.ThreadHelper;
 
 namespace Microsoft.VisualStudio.IntegrationTest.Utilities.InProcess
 {
@@ -264,25 +263,28 @@ public bool IsCaretOnScreen()
 
         public ClassifiedToken[] GetLightbulbPreviewClassifications(string menuText)
         {
-            return ExecuteOnActiveView(view =>
+            return ThreadHelper.JoinableTaskFactory.Run(async () =>
             {
+                await ThreadHelper.JoinableTaskFactory.SwitchToMainThreadAsync();
+
+                var view = GetActiveTextView();
                 var broker = GetComponentModel().GetService<ILightBulbBroker>();
                 var classifierAggregatorService = GetComponentModelService<IViewClassifierAggregatorService>();
-                return GetLightbulbPreviewClassifications(
+                return await GetLightbulbPreviewClassificationsAsync(
                     menuText,
                     broker,
                     view,
-                    classifierAggregatorService);
+                    classifierAggregatorService).ConfigureAwait(false);
             });
         }
 
-        private ClassifiedToken[] GetLightbulbPreviewClassifications(
-                    string menuText,
-                    ILightBulbBroker broker,
-                    IWpfTextView view,
-                    IViewClassifierAggregatorService viewClassifierAggregator)
+        private async Task<ClassifiedToken[]> GetLightbulbPreviewClassificationsAsync(
+            string menuText,
+            ILightBulbBroker broker,
+            IWpfTextView view,
+            IViewClassifierAggregatorService viewClassifierAggregator)
         {
-            LightBulbHelper.WaitForLightBulbSession(broker, view);
+            await LightBulbHelper.WaitForLightBulbSessionAsync(broker, view).ConfigureAwait(true);
 
             var bufferType = view.TextBuffer.ContentType.DisplayName;
             if (!broker.IsLightBulbSessionActive(view))
@@ -311,7 +313,7 @@ public ClassifiedToken[] GetLightbulbPreviewClassifications(string menuText)
                 }
 
                 IWpfTextView preview = null;
-                object pane = HostWaitHelper.PumpingWaitResult(set.GetPreviewAsync(CancellationToken.None));
+                object pane = await set.GetPreviewAsync(CancellationToken.None).ConfigureAwait(true);
                 if (pane is System.Windows.Controls.UserControl)
                 {
                     var container = ((System.Windows.Controls.UserControl)pane).FindName("PreviewDockPanel") as DockPanel;

src/VisualStudio/IntegrationTest/TestUtilities/InProcess/TextViewWindow_InProc.cs

@@ -5,6 +5,7 @@
 using System.Linq;
 using System.Text;
 using System.Threading;
+using System.Threading.Tasks;
 using Microsoft.CodeAnalysis.CodeFixes;
 using Microsoft.CodeAnalysis.Editor.Implementation.Suggestions;
 using Microsoft.VisualStudio.Language.Intellisense;
@@ -15,6 +16,7 @@
 using Microsoft.VisualStudio.Text.Operations;
 using Microsoft.VisualStudio.Text.Tagging;
 using OLECMDEXECOPT = Microsoft.VisualStudio.OLE.Interop.OLECMDEXECOPT;
+using ThreadHelper = Microsoft.VisualStudio.Shell.ThreadHelper;
 
 namespace Microsoft.VisualStudio.IntegrationTest.Utilities.InProcess
 {
@@ -75,11 +77,16 @@ public void ShowLightBulb()
         }
 
         public void WaitForLightBulbSession()
-            => ExecuteOnActiveView(view =>
+        {
+            ThreadHelper.JoinableTaskFactory.Run(async () =>
             {
+                await ThreadHelper.JoinableTaskFactory.SwitchToMainThreadAsync();
+
+                var view = GetActiveTextView();
                 var broker = GetComponentModel().GetService<ILightBulbBroker>();
-                LightBulbHelper.WaitForLightBulbSession(broker, view);
+                await LightBulbHelper.WaitForLightBulbSessionAsync(broker, view).ConfigureAwait(false);
             });
+        }
 
         /// <remarks>
         /// This method does not wait for async operations before
@@ -274,13 +281,18 @@ public bool IsLightBulbSessionExpanded()
        });
 
         public string[] GetLightBulbActions()
-            => ExecuteOnActiveView(view =>
+        {
+            return ThreadHelper.JoinableTaskFactory.Run(async () =>
             {
+                await ThreadHelper.JoinableTaskFactory.SwitchToMainThreadAsync();
+
+                var view = GetActiveTextView();
                 var broker = GetComponentModel().GetService<ILightBulbBroker>();
-                return GetLightBulbActions(broker, view).Select(a => a.DisplayText).ToArray();
+                return (await GetLightBulbActionsAsync(broker, view)).Select(a => a.DisplayText).ToArray();
             });
+        }
 
-        private IEnumerable<ISuggestedAction> GetLightBulbActions(ILightBulbBroker broker, IWpfTextView view)
+        private async Task<IEnumerable<ISuggestedAction>> GetLightBulbActionsAsync(ILightBulbBroker broker, IWpfTextView view)
         {
             if (!broker.IsLightBulbSessionActive(view))
             {
@@ -300,19 +312,23 @@ private IEnumerable<ISuggestedAction> GetLightBulbActions(ILightBulbBroker broke
                 actionSets = Array.Empty<SuggestedActionSet>();
             }
 
-            return SelectActions(actionSets);
+            return await SelectActionsAsync(actionSets);
         }
 
         public void ApplyLightBulbAction(string actionName, FixAllScope? fixAllScope, bool blockUntilComplete)
         {
             var lightBulbAction = GetLightBulbApplicationAction(actionName, fixAllScope);
-            if (blockUntilComplete)
+            var task = ThreadHelper.JoinableTaskFactory.RunAsync(async () =>
             {
-                ExecuteOnActiveView(lightBulbAction);
-            }
-            else
+                await ThreadHelper.JoinableTaskFactory.SwitchToMainThreadAsync();
+
+                var activeTextView = GetActiveTextView();
+                await lightBulbAction(activeTextView);
+            });
+
+            if (blockUntilComplete)
             {
-                BeginInvokeExecuteOnActiveView(lightBulbAction);
+                task.Join();
             }
         }
 
@@ -322,13 +338,13 @@ public void ApplyLightBulbAction(string actionName, FixAllScope? fixAllScope, bo
         private void BeginInvokeExecuteOnActiveView(Action<IWpfTextView> action)
             => BeginInvokeOnUIThread(GetExecuteOnActionViewCallback(action));
 
-        private Action<IWpfTextView> GetLightBulbApplicationAction(string actionName, FixAllScope? fixAllScope)
+        private Func<IWpfTextView, Task> GetLightBulbApplicationAction(string actionName, FixAllScope? fixAllScope)
         {
-            return view =>
+            return async view =>
             {
                 var broker = GetComponentModel().GetService<ILightBulbBroker>();
 
-                var actions = GetLightBulbActions(broker, view).ToArray();
+                var actions = (await GetLightBulbActionsAsync(broker, view).ConfigureAwait(true)).ToArray();
                 var action = actions.FirstOrDefault(a => a.DisplayText == actionName);
 
                 if (action == null)
@@ -351,8 +367,8 @@ private Action<IWpfTextView> GetLightBulbApplicationAction(string actionName, Fi
                         throw new InvalidOperationException($"Suggested action '{action.DisplayText}' does not support FixAllOccurrences.");
                     }
 
-                    var actionSetsForAction = HostWaitHelper.PumpingWaitResult(action.GetActionSetsAsync(CancellationToken.None));
-                    action = GetFixAllSuggestedAction(actionSetsForAction, fixAllScope.Value);
+                    var actionSetsForAction = await action.GetActionSetsAsync(CancellationToken.None).ConfigureAwait(true);
+                    action = await GetFixAllSuggestedActionAsync(actionSetsForAction, fixAllScope.Value).ConfigureAwait(true);
                     if (action == null)
                     {
                         throw new InvalidOperationException($"Unable to find FixAll in {fixAllScope.ToString()} code fix for suggested action '{action.DisplayText}'.");
@@ -371,7 +387,7 @@ private Action<IWpfTextView> GetLightBulbApplicationAction(string actionName, Fi
             };
         }
 
-        private IEnumerable<ISuggestedAction> SelectActions(IEnumerable<SuggestedActionSet> actionSets)
+        private async Task<IEnumerable<ISuggestedAction>> SelectActionsAsync(IEnumerable<SuggestedActionSet> actionSets)
         {
             var actions = new List<ISuggestedAction>();
 
@@ -384,7 +400,7 @@ private IEnumerable<ISuggestedAction> SelectActions(IEnumerable<SuggestedActionS
                         foreach (var action in actionSet.Actions)
                         {
                             actions.Add(action);
-                            actions.AddRange(SelectActions(HostWaitHelper.PumpingWaitResult(action.GetActionSetsAsync(CancellationToken.None))));
+                            actions.AddRange(await SelectActionsAsync(await action.GetActionSetsAsync(CancellationToken.None)));
                         }
                     }
                 }
@@ -393,7 +409,7 @@ private IEnumerable<ISuggestedAction> SelectActions(IEnumerable<SuggestedActionS
             return actions;
         }
 
-        private static FixAllSuggestedAction GetFixAllSuggestedAction(IEnumerable<SuggestedActionSet> actionSets, FixAllScope fixAllScope)
+        private static async Task<FixAllSuggestedAction> GetFixAllSuggestedActionAsync(IEnumerable<SuggestedActionSet> actionSets, FixAllScope fixAllScope)
         {
             foreach (var actionSet in actionSets)
             {
@@ -410,8 +426,8 @@ private static FixAllSuggestedAction GetFixAllSuggestedAction(IEnumerable<Sugges
 
                     if (action.HasActionSets)
                     {
-                        var nestedActionSets = HostWaitHelper.PumpingWaitResult(action.GetActionSetsAsync(CancellationToken.None));
-                        fixAllSuggestedAction = GetFixAllSuggestedAction(nestedActionSets, fixAllScope);
+                        var nestedActionSets = await action.GetActionSetsAsync(CancellationToken.None);
+                        fixAllSuggestedAction = await GetFixAllSuggestedActionAsync(nestedActionSets, fixAllScope);
                         if (fixAllSuggestedAction != null)
                         {
                             return fixAllSuggestedAction;

src/VisualStudio/IntegrationTest/TestUtilities/LightBulbHelper.cs

@@ -3,23 +3,27 @@
 using System;
 using System.Threading.Tasks;
 using Microsoft.VisualStudio.Language.Intellisense;
+using Microsoft.VisualStudio.Text.Editor;
 
 namespace Microsoft.VisualStudio.IntegrationTest.Utilities
 {
     public static class LightBulbHelper
     {
-        public static bool WaitForLightBulbSession(ILightBulbBroker broker, Microsoft.VisualStudio.Text.Editor.IWpfTextView view)
-            => Helper.Retry(() => {
+        public static Task<bool> WaitForLightBulbSessionAsync(ILightBulbBroker broker, IWpfTextView view)
+        {
+            return Helper.RetryAsync(async () =>
+            {
                 if (broker.IsLightBulbSessionActive(view))
                 {
                     return true;
                 }
 
-                // checking whether there is any suggested action is async up to editor layer and our waiter doesnt track up to that point.
+                // checking whether there is any suggested action is async up to editor layer and our waiter doesn't track up to that point.
                 // so here, we have no other way than sleep (with timeout) to see LB is available.
-                HostWaitHelper.PumpingWait(Task.Delay(TimeSpan.FromSeconds(1)));
+                await Task.Delay(TimeSpan.FromSeconds(1)).ConfigureAwait(true);
 
                 return broker.IsLightBulbSessionActive(view);
-            }, TimeSpan.FromSeconds(0));
+            }, TimeSpan.Zero);
+        }
     }
 }