Changed how the test awaits on executed queries to make it more relia…

…ble (#427)
microsoft · Aug 4, 2017 · 70d658b · 70d658b
1 parent aa725c5
commit 70d658b
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 12 deletions.
diff --git a/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/ResultOnlyContext.cs b/src/Microsoft.SqlTools.ServiceLayer/QueryExecution/ResultOnlyContext.cs
@@ -15,10 +15,10 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
     /// </summary>
     public class ResultOnlyContext<TResult> : IEventSender
     {
-        private readonly RequestContext<TResult> origContext;
+        private readonly RequestContext<TResult> OrigContext;
 
         public ResultOnlyContext(RequestContext<TResult> context) {
-            origContext = context;
+            OrigContext = context;
         }
 
         public virtual Task SendEvent<TParams>(EventType<TParams> eventType, TParams eventParams)
@@ -27,5 +27,10 @@ public virtual Task SendEvent<TParams>(EventType<TParams> eventType, TParams eve
             // in the future this could be used to roll up events and send them back in the result
             return Task.FromResult(true);
         }
+
+        public virtual Task SendError(string errorMessage, int errorCode = 0)
+        {
+            return OrigContext.SendError(errorMessage, errorCode);
+        }
     }
 }
diff --git a/...osoft.SqlTools.ServiceLayer.UnitTests/QueryExecution/Execution/ServiceIntegrationTests.cs b/...osoft.SqlTools.ServiceLayer.UnitTests/QueryExecution/Execution/ServiceIntegrationTests.cs
@@ -486,16 +486,10 @@ public async Task SimpleExecuteMultipleQueriesTest()
 
             await Task.WhenAll(queryService.ActiveSimpleExecuteRequests.Values);
 
-            var queries = queryService.ActiveQueries.Values.Take(2).ToArray();
-            Query q1 = queries[0];
-            Query q2 = queries[1];
-
-            Assert.NotNull(q1);
-            Assert.NotNull(q2);
-
-            // wait on the task to finish
-            q1.ExecutionTask.Wait();
-            q2.ExecutionTask.Wait();
+            var queries = queryService.ActiveQueries.Values.ToArray();
+            var queryTasks = queries.Select(query => query.ExecutionTask);
+
+            await Task.WhenAll(queryTasks);
 
             efv1.Validate();
             efv2.Validate();