Fixes tests to use the correct context to enable async refactor

src/Microsoft.OData.Cli/ODataCliFileHandler.cs

@@ -91,16 +91,16 @@ public Task<string> AddFileAsync(string fileName, string targetPath, ODataFileOp
         /// Emits the container property attribute. The lastest version gets installed. so we'll always emit this property.
         /// </summary>
         /// <returns>A bool indicating whether to emit the container property or not</returns>
-        public bool EmitContainerPropertyAttribute()
+        public Task<bool> EmitContainerPropertyAttributeAsync()
         {
-            return true;
+            return Task.FromResult(true);
         }
 
         /// <summary>
         /// Sets the CSDL file as an embedded resource
         /// </summary>
         /// <param name="fileName">The name of the file to set as embedded resource</param>
-        public void SetFileAsEmbeddedResource(string fileName)
+        public Task SetFileAsEmbeddedResourceAsync(string fileName)
         {
             if (this.project != null)
             {
@@ -111,6 +111,8 @@ public void SetFileAsEmbeddedResource(string fileName)
                     ProjectHelper.AddProjectItem(this.project, Constants.EmbeddedResourceTag, fileName);
                 }
             }
+
+            return Task.CompletedTask;
         }
     }
 }

src/Microsoft.OData.CodeGen/CodeGeneration/V4CodeGenDescriptor.cs

@@ -35,13 +35,13 @@ public V4CodeGenDescriptor(IFileHandler fileHandler, IMessageLogger logger, IPac
 
         public override async Task AddNugetPackagesAsync()
         {
-            await MessageLogger.WriteMessageAsync(LogMessageCategory.Information, "Adding Nuget Packages...").ConfigureAwait(false);
+            await MessageLogger.WriteMessageAsync(LogMessageCategory.Information, "Adding Nuget Packages...");
 
 
             foreach (var nugetPackage in Common.Constants.V4NuGetPackages)
-                await PackageInstaller.CheckAndInstallNuGetPackageAsync(Common.Constants.NuGetOnlineRepository, nugetPackage).ConfigureAwait(false);
+                await PackageInstaller.CheckAndInstallNuGetPackageAsync(Common.Constants.NuGetOnlineRepository, nugetPackage);
 
-            await MessageLogger.WriteMessageAsync(LogMessageCategory.Information, "Nuget Packages were installed.").ConfigureAwait(false);
+            await MessageLogger.WriteMessageAsync(LogMessageCategory.Information, "Nuget Packages were installed.");
         }
 
         public override async Task AddGeneratedClientCodeAsync(string metadata, string outputDirectory, LanguageOption languageOption, ServiceConfiguration serviceConfiguration)
@@ -96,7 +96,7 @@ private async Task AddT4FileAsync(string metadata, string outputDirectory, Langu
 
             await FileHandler.AddFileAsync(csdlTempFile, metadataFile, new ODataFileOptions { SuppressOverwritePrompt = true });
 
-            FileHandler.SetFileAsEmbeddedResource(csdlFileName);
+            await FileHandler.SetFileAsEmbeddedResourceAsync(csdlFileName);
 
             var t4TempFile = Path.GetTempFileName();
 
@@ -193,8 +193,8 @@ private async Task AddGeneratedCodeAsync(string metadata, string outputDirectory
             var metadataFile = Path.Combine(referenceFolder, csdlFileName);
             await FileHandler.AddFileAsync(tempFile, metadataFile, new ODataFileOptions { SuppressOverwritePrompt = true });
 
-            FileHandler.SetFileAsEmbeddedResource(csdlFileName);
-            t4CodeGenerator.EmitContainerPropertyAttribute = FileHandler.EmitContainerPropertyAttribute();
+            await FileHandler.SetFileAsEmbeddedResourceAsync(csdlFileName);
+            t4CodeGenerator.EmitContainerPropertyAttribute = await FileHandler.EmitContainerPropertyAttributeAsync();
 
             t4CodeGenerator.MetadataFilePath = metadataFile;
             t4CodeGenerator.MetadataFileRelativePath = csdlFileName;
@@ -207,7 +207,7 @@ private async Task AddGeneratedCodeAsync(string metadata, string outputDirectory
                 {
                     foreach (var err in t4CodeGenerator.Errors)
                     {
-                        await MessageLogger.WriteMessageAsync(LogMessageCategory.Warning, err.ToString()).ConfigureAwait(false);
+                        await MessageLogger.WriteMessageAsync(LogMessageCategory.Warning, err.ToString());
                     }
                 }
             }

src/Microsoft.OData.CodeGen/FileHandling/IFileHandler.cs

@@ -26,12 +26,12 @@ public interface IFileHandler
         /// Sets a file as an embedded resource
         /// </summary>
         /// <param name="fileName">The name of the file to set as an embedded resource</param>
-        void SetFileAsEmbeddedResource(string fileName);
+        Task SetFileAsEmbeddedResourceAsync(string fileName);
 
         /// <summary>
         /// Emits container property attribute
         /// </summary>
         /// <returns></returns>
-        bool EmitContainerPropertyAttribute();
+        Task<bool> EmitContainerPropertyAttributeAsync();
     }
 }

src/Microsoft.OData.CodeGen/Microsoft.OData.CodeGen.csproj

@@ -79,6 +79,7 @@
     <Compile Include="Common\MetadataReader.cs" />
     <Compile Include="FileHandling\IFileHandler.cs" />
     <Compile Include="FileHandling\ODataFileOptions.cs" />
+    <Compile Include="Threading\IThreadHelper.cs" />
     <Compile Include="Logging\IMessageLogger.cs" />
     <Compile Include="Logging\LogMessageCategory.cs" />
     <Compile Include="Models\BaseUserSettings.cs" />

src/Microsoft.OData.CodeGen/Threading/IThreadHelper.cs

@@ -0,0 +1,34 @@
+//-----------------------------------------------------------------------------
+// <copyright file="IThreadHelper.cs" company=".NET Foundation">
+//      Copyright (c) .NET Foundation and Contributors. All rights reserved.
+//      See License.txt in the project root for license information.
+// </copyright>
+//----------------------------------------------------------------------------
+
+using System;
+using System.Threading.Tasks;
+
+namespace Microsoft.OData.CodeGen
+{
+    /// <summary>
+    /// A Thread helper to assist users to marshal certain work in certain threads.
+    /// </summary>
+    public interface IThreadHelper
+    {
+        /// <summary>
+        /// Runs the block provided in the background.
+        /// </summary>
+        /// <typeparam name="T">Return type generic type parameter.</typeparam>
+        /// <param name="backgroundTask">Function containing the task to be ran in the background thread executor.</param>
+        /// <returns>A task respresenting the completion of the task.</returns>
+        Task<T> RunAsync<T>(Func<Task<T>> backgroundTask);
+
+        /// <summary>
+        /// Runs the block provided in the foreground UI thread.
+        /// </summary>
+        /// <typeparam name="T">Return type generic type parameter.</typeparam>
+        /// <param name="foregroundTask">A func representing the foreground task.</param>
+        /// <returns>A task respresenting the completion of the task.</returns>
+        Task<T> RunInUiThreadAsync<T>(Func<T> foregroundTask);
+    }
+}

src/ODataConnectedService.Shared/ConnectedServiceFileHandler.cs

@@ -11,7 +11,10 @@
 using Microsoft.OData.CodeGen;
 using Microsoft.OData.CodeGen.FileHandling;
 using Microsoft.VisualStudio.ConnectedServices;
+using Microsoft.VisualStudio.Shell;
+using Microsoft.VisualStudio.Threading;
 using VSLangProj;
+using Task = System.Threading.Tasks.Task;
 
 namespace Microsoft.OData.ConnectedService
 {
@@ -21,17 +24,20 @@ namespace Microsoft.OData.ConnectedService
     public class ConnectedServiceFileHandler : IFileHandler
     {
         private ConnectedServiceHandlerContext Context;
+        private readonly IThreadHelper threadHelper;
+
         public Project Project { get; private set; }
 
         /// <summary>
         /// Creates an instance of <see cref="ConnectedServiceFileHandler"/>
         /// </summary>
         /// <param name="context">The <see cref="ConnectedServiceHandlerContext"/ object></param>
         /// <param name="project">An object of the project.</param>
-        public ConnectedServiceFileHandler(ConnectedServiceHandlerContext context, Project project )
+        public ConnectedServiceFileHandler(ConnectedServiceHandlerContext context, Project project, IThreadHelper threadHelper)
         {
             this.Context = context;
             this.Project = project;
+            this.threadHelper = threadHelper;
         }
 
         /// <summary>
@@ -57,39 +63,49 @@ public async Task<string> AddFileAsync(string fileName, string targetPath, OData
         /// Sets the CSDL file as an embedded resource
         /// </summary>
         /// <param name="fileName">The name of the file to set as embedded resource</param>
-        public void SetFileAsEmbeddedResource(string fileName)
+        public async Task SetFileAsEmbeddedResourceAsync(string fileName)
         {
-            var dte = Microsoft.VisualStudio.Shell.Package.GetGlobalService(typeof(DTE)) as DTE;
-            if (dte != null)
+            await this.threadHelper.RunInUiThreadAsync(() =>
             {
-                var projectItem = this.Project.ProjectItems.Item("Connected Services").ProjectItems.Item(((ODataConnectedServiceInstance)this.Context.ServiceInstance).ServiceConfig.ServiceName).ProjectItems.Item(fileName);
-                projectItem.Properties.Item("BuildAction").Value = prjBuildAction.prjBuildActionEmbeddedResource;
-            }
+#pragma warning disable VSTHRD010 // Invoke single-threaded types on Main thread
+                if (Package.GetGlobalService(typeof(DTE)) is DTE dte)
+                {
+                    var projectItem = this.Project.ProjectItems.Item("Connected Services").ProjectItems.Item(((ODataConnectedServiceInstance)this.Context.ServiceInstance).ServiceConfig.ServiceName).ProjectItems.Item(fileName);
+                    projectItem.Properties.Item("BuildAction").Value = prjBuildAction.prjBuildActionEmbeddedResource;
+                    return true;
+                }
+#pragma warning restore VSTHRD010 // Invoke single-threaded types on Main thread
+                return false;
+            }).ConfigureAwait(true);
         }
 
         /// <summary>
         /// Sets the container property attribute to either true or false
         /// </summary>
         /// <returns>A value of either true or false</returns>
-        public bool EmitContainerPropertyAttribute()
+        public async Task<bool> EmitContainerPropertyAttributeAsync()
         {
-            var vsProject = this.Project.Object as VSProject;
-            if (vsProject != null)
+           return await threadHelper.RunInUiThreadAsync(() =>
             {
-                foreach (Reference reference in vsProject.References)
+#pragma warning disable VSTHRD010 // Invoke single-threaded types on Main thread
+                if (this.Project.Object is VSProject vsProject)
                 {
-                    if (reference.SourceProject == null)
+                    foreach (Reference reference in vsProject.References)
                     {
-                        // Assembly reference (For project reference, SourceProject != null)
-                        if (reference.Name.Equals("Microsoft.OData.Client", StringComparison.Ordinal))
+                        if (reference.SourceProject == null)
                         {
-                            return Version.Parse(reference.Version) > Version.Parse("7.6.4.0");
+                            // Assembly reference (For project reference, SourceProject != null)
+                            if (reference.Name.Equals("Microsoft.OData.Client", StringComparison.Ordinal))
+                            {
+                                return Version.Parse(reference.Version) > Version.Parse("7.6.4.0");
+                            }
                         }
                     }
                 }
-            }
+#pragma warning restore VSTHRD010 // Invoke single-threaded types on Main thread
 
-            return false;
+                return false;
+            }).ConfigureAwait(true);
         }
     }
 }

src/ODataConnectedService.Shared/ConnectedServicePackageInstaller.cs

@@ -26,6 +26,7 @@ public class ConnectedServicePackageInstaller : IPackageInstaller
         public Project Project { get; private set; }
         public IMessageLogger MessageLogger { get; private set; }
         public IVsPackageInstaller PackageInstaller { get; protected set; }
+
         public IVsPackageInstallerServices PackageInstallerServices { get; protected set; }
 
         /// <summary>
@@ -68,8 +69,7 @@ public async Task CheckAndInstallNuGetPackageAsync(string packageSource, string
                 {
                     if (!PackageInstallerServices.IsPackageInstalled(this.Project, packageName))
                     {
-                        Version packageVersion = null;
-                        PackageInstaller.InstallPackage(packageSource, this.Project, packageName, packageVersion, false);
+                        PackageInstaller.InstallPackage(packageSource, this.Project, packageName, (string)null, false); ;
 
                         await (this.MessageLogger?.WriteMessageAsync(LogMessageCategory.Information, $"Nuget Package \"{packageName}\" for OData client was added.")).ConfigureAwait(false);
                     }

src/ODataConnectedService.Shared/ConnectedServiceThreadHelper.cs

@@ -0,0 +1,24 @@
+using System;
+
+using System.Threading.Tasks;
+using Microsoft.OData.CodeGen;
+using Microsoft.VisualStudio.Shell;
+
+namespace ODataConnectedService.Shared
+{
+    internal class ConnectedServiceThreadHelper : IThreadHelper
+    {
+        /// <inheritdoc/>
+        public async Task<T> RunAsync<T>(Func<Task<T>> backgroundTask)
+        {
+            return await ThreadHelper.JoinableTaskFactory.RunAsync(backgroundTask);
+        }
+
+        /// <inheritdoc/>
+        public async Task<T> RunInUiThreadAsync<T>(Func<T> foregroundTask)
+        {
+            await ThreadHelper.JoinableTaskFactory.SwitchToMainThreadAsync();
+            return foregroundTask();
+        }
+    }
+}

src/ODataConnectedService.Shared/ODataConnectedService.Shared.projitems

@@ -17,6 +17,7 @@
     <Compile Include="$(MSBuildThisFileDirectory)ConnectedServiceFileHandler.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)ConnectedServiceMessageLogger.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)ConnectedServicePackageInstaller.cs" />
+    <Compile Include="$(MSBuildThisFileDirectory)ConnectedServiceThreadHelper.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Converters\BoolToVisibilityConverter.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)Converters\SchemaTypeModelToVisibilityConverter.cs" />
     <Compile Include="$(MSBuildThisFileDirectory)GlobalSuppressions.cs" />

src/ODataConnectedService.Shared/ODataConnectedServiceHandler.cs

@@ -11,9 +11,11 @@
 using System.Threading;
 using System.Threading.Tasks;
 using EnvDTE;
+using Microsoft.OData.CodeGen;
 using Microsoft.OData.CodeGen.CodeGeneration;
 using Microsoft.OData.ConnectedService.Common;
 using Microsoft.VisualStudio.ConnectedServices;
+using ODataConnectedService.Shared;
 
 namespace Microsoft.OData.ConnectedService
 {
@@ -49,7 +51,6 @@ private async Task<BaseCodeGenDescriptor> SaveServiceInstanceAsync(ConnectedServ
         {
             Project project = ProjectHelper.GetProjectFromHierarchy(context.ProjectHierarchy);
             var serviceInstance = (ODataConnectedServiceInstance)context.ServiceInstance;
-
             var codeGenDescriptor = await GenerateCodeAsync(serviceInstance.ServiceConfig.Endpoint, serviceInstance.ServiceConfig.EdmxVersion, context, project).ConfigureAwait(false);            
             if (!serviceInstance.ServiceConfig.StoreCustomHttpHeaders)
             {
@@ -89,7 +90,9 @@ private async Task<BaseCodeGenDescriptor> GenerateCodeAsync(string metadataUri,
                 languageOption = LanguageOption.GenerateCSharpCode;
             }
 
-            BaseCodeGenDescriptor codeGenDescriptor = codeGenDescriptorFactory.Create(edmxVersion, new ConnectedServiceFileHandler(context, project), new ConnectedServiceMessageLogger(context), new ConnectedServicePackageInstaller(context, project, new ConnectedServiceMessageLogger(context)));
+            var vsThreadHelper = new ConnectedServiceThreadHelper();
+
+            BaseCodeGenDescriptor codeGenDescriptor = codeGenDescriptorFactory.Create(edmxVersion, new ConnectedServiceFileHandler(context, project, vsThreadHelper), new ConnectedServiceMessageLogger(context), new ConnectedServicePackageInstaller(context, project, new ConnectedServiceMessageLogger(context)));
             await codeGenDescriptor.AddNugetPackagesAsync().ConfigureAwait(false);
             await codeGenDescriptor.AddGeneratedClientCodeAsync(metadataUri, outputDirectory, (Microsoft.OData.CodeGen.Common.LanguageOption)languageOption, ((ODataConnectedServiceInstance)context.ServiceInstance).ServiceConfig).ConfigureAwait(false);
             return codeGenDescriptor;