Feature: Progressive Messages (#208)

This change is a reworking of the way that messages are sent to clients from the service layer. It is also a reworking of the protocol to ensure that all formulations of query send back events to the client in a deterministic ordering. To support the first change:
* Added a new event that will be sent when a message is generated
* Messages now indicate which Batch (if any) generated them
* Messages now indicate if they were error level
* Removed message storage in Batch objects and BatchSummary objects
* Batch objects no longer have error state

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Batch.cs

@@ -40,16 +40,16 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         private DateTime executionStartTime;
 
+        /// <summary>
+        /// Whether or not any messages have been sent
+        /// </summary>
+        private bool messagesSent;
+
         /// <summary>
         /// Factory for creating readers/writers for the output of the batch
         /// </summary>
         private readonly IFileStreamFactory outputFileFactory;
 
-        /// <summary>
-        /// Internal representation of the messages so we can modify internally
-        /// </summary>
-        internal readonly List<ResultMessage> resultMessages;
-
         /// <summary>
         /// Internal representation of the result sets so we can modify internally
         /// </summary>
@@ -71,7 +71,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             HasExecuted = false;
             Id = ordinalId;
             resultSets = new List<ResultSet>();
-            resultMessages = new List<ResultMessage>();
             this.outputFileFactory = outputFileFactory;
         }
 
@@ -83,11 +82,22 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// <param name="batch">The batch that completed</param>
         public delegate Task BatchAsyncEventHandler(Batch batch);
 
+        /// <summary>
+        /// Asynchronous handler for when a message is emitted by the sql connection
+        /// </summary>
+        /// <param name="message">The message that was emitted</param>
+        public delegate Task BatchAsyncMessageHandler(ResultMessage message);
+
         /// <summary>
         /// Event that will be called when the batch has completed execution
         /// </summary>
         public event BatchAsyncEventHandler BatchCompletion;
 
+        /// <summary>
+        /// Event that will be called when a message has been emitted
+        /// </summary>
+        public event BatchAsyncMessageHandler BatchMessageSent;
+
         /// <summary>
         /// Event to call when the batch has started execution
         /// </summary>
@@ -132,11 +142,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         public string ExecutionStartTimeStamp { get { return executionStartTime.ToString("o"); } }
 
-        /// <summary>
-        /// Whether or not this batch has an error
-        /// </summary>
-        public bool HasError { get; set; }
-
         /// <summary>
         /// Whether or not this batch has been executed, regardless of success or failure 
         /// </summary>
@@ -147,14 +152,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         public int Id { get; private set; }
 
-        /// <summary>
-        /// Messages that have come back from the server
-        /// </summary>
-        public IEnumerable<ResultMessage> ResultMessages
-        {
-            get { return resultMessages; }
-        }
-
         /// <summary>
         /// The result sets of the batch execution
         /// </summary>
@@ -187,7 +184,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 // Batch summary with information available at start
                 BatchSummary summary = new BatchSummary
                 {
-                    HasError = HasError,
                     Id = Id,
                     Selection = Selection,
                     ExecutionStart = ExecutionStartTimeStamp
@@ -197,7 +193,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 if (HasExecuted)
                 {
                     summary.ResultSetSummaries = ResultSummaries;
-                    summary.Messages = ResultMessages.ToArray();
                     summary.ExecutionEnd = ExecutionEndTimeStamp;
                     summary.ExecutionElapsed = ExecutionElapsedTime;
                 }
@@ -244,7 +239,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 {
                     // Register the message listener to *this instance* of the batch
                     // Note: This is being done to associate messages with batches
-                    sqlConn.GetUnderlyingConnection().InfoMessage += StoreDbMessage;
+                    sqlConn.GetUnderlyingConnection().InfoMessage += ServerMessageHandler;
                     command = sqlConn.GetUnderlyingConnection().CreateCommand();
 
                     // Add a handler for when the command completes
@@ -298,27 +293,25 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 
                         // If there were no messages, for whatever reason (NO COUNT set, messages 
                         // were emitted, records returned), output a "successful" message
-                        if (resultMessages.Count == 0)
+                        if (!messagesSent)
                         {
-                            resultMessages.Add(new ResultMessage(SR.QueryServiceCompletedSuccessfully));
+                            await SendMessage(SR.QueryServiceCompletedSuccessfully, false);
                         }
                     }
                 }
             }
             catch (DbException dbe)
             {
-                HasError = true;
+                await UnwrapDbException(dbe);
-                UnwrapDbException(dbe);
             }
             catch (TaskCanceledException)
             {
-                resultMessages.Add(new ResultMessage(SR.QueryServiceQueryCancelled));
+                await SendMessage(SR.QueryServiceQueryCancelled, false);
                 throw;
             }
             catch (Exception e)
             {
-                HasError = true;
+                await SendMessage(SR.QueryServiceQueryFailed(e.Message), true);
-                resultMessages.Add(new ResultMessage(SR.QueryServiceQueryFailed(e.Message)));
                 throw;
             }
             finally
@@ -327,7 +320,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 ReliableSqlConnection sqlConn = conn as ReliableSqlConnection;
                 if (sqlConn != null)
                 {
-                    sqlConn.GetUnderlyingConnection().InfoMessage -= StoreDbMessage;
+                    sqlConn.GetUnderlyingConnection().InfoMessage -= ServerMessageHandler;
                 }
 
                 // Mark that we have executed
@@ -400,6 +393,19 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
 
         #region Private Helpers
 
+        private async Task SendMessage(string message, bool isError)
+        {
+            // If the message event is null, this is a no-op
+            if (BatchMessageSent == null)
+            {
+                return;
+            }
+
+            // State that we've sent any message, and send it
+            messagesSent = true;
+            await BatchMessageSent(new ResultMessage(message, isError, Id));
+        }
+
         /// <summary>
         /// Handler for when the StatementCompleted event is fired for this batch's command. This
         /// will be executed ONLY when there is a rowcount to report. If this event is not fired
@@ -410,17 +416,10 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         internal void StatementCompletedHandler(object sender, StatementCompletedEventArgs args)
         {
             // Add a message for the number of rows the query returned
-            string message;
+            string message = args.RecordCount == 1
-            if (args.RecordCount == 1)
+                ? SR.QueryServiceAffectedOneRow
-            {
+                : SR.QueryServiceAffectedRows(args.RecordCount);
-                message = SR.QueryServiceAffectedOneRow;
+            SendMessage(message, false).Wait();
-            }
-            else
-            {
-                message = SR.QueryServiceAffectedRows(args.RecordCount);
-            }
-
-            resultMessages.Add(new ResultMessage(message));
         }
 
         /// <summary>
@@ -430,30 +429,30 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         /// <param name="sender">Object that fired the event</param>
         /// <param name="args">Arguments from the event</param>
-        private void StoreDbMessage(object sender, SqlInfoMessageEventArgs args)
+        private void ServerMessageHandler(object sender, SqlInfoMessageEventArgs args)
         {
-            resultMessages.Add(new ResultMessage(args.Message));
+            SendMessage(args.Message, false).Wait();
         }
 
         /// <summary>
-        /// Attempts to convert a <see cref="DbException"/> to a <see cref="SqlException"/> that
+        /// Attempts to convert an <see cref="Exception"/> to a <see cref="SqlException"/> that
         /// contains much more info about Sql Server errors. The exception is then unwrapped and
-        /// messages are formatted and stored in <see cref="ResultMessages"/>. If the exception
+        /// messages are formatted and sent to the extension. If the exception cannot be 
-        /// cannot be converted to SqlException, the message is written to the messages list.
+        /// converted to SqlException, the message is written to the messages list.
         /// </summary>
         /// <param name="dbe">The exception to unwrap</param>
-        internal void UnwrapDbException(DbException dbe)
+        private async Task UnwrapDbException(Exception dbe)
         {
             SqlException se = dbe as SqlException;
             if (se != null)
             {
                 var errors = se.Errors.Cast<SqlError>().ToList();
+                
                 // Detect user cancellation errors
                 if (errors.Any(error => error.Class == 11 && error.Number == 0))
                 {
                     // User cancellation error, add the single message
-                    HasError = false;
+                    await SendMessage(SR.QueryServiceQueryCancelled, false);
-                    resultMessages.Add(new ResultMessage(SR.QueryServiceQueryCancelled));
                 }
                 else
                 {
@@ -464,13 +463,13 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                         string message = string.Format("Msg {0}, Level {1}, State {2}, Line {3}{4}{5}",
                             error.Number, error.Class, error.State, lineNumber,
                             Environment.NewLine, error.Message);
-                        resultMessages.Add(new ResultMessage(message));
+                        await SendMessage(message, true);
                     }
                 }
             }
             else
             {
-                resultMessages.Add(new ResultMessage(dbe.Message));
+                await SendMessage(dbe.Message, true);
             }
         }
 

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/BatchSummary.cs

@@ -25,11 +25,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
         /// </summary>
         public string ExecutionStart { get; set; }
 
-        /// <summary>
-        /// Whether or not the batch was successful. True indicates errors, false indicates success
-        /// </summary>
-        public bool HasError { get; set; }
-
         /// <summary>
         /// The ID of the result set within the query results
         /// </summary>
@@ -40,11 +35,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
         /// </summary>
         public SelectionData Selection { get; set; }
 
-        /// <summary>
-        /// Any messages that came back from the server during execution of the batch
-        /// </summary>
-        public ResultMessage[] Messages { get; set; }
-
         /// <summary>
         /// The summaries of the result sets inside the batch
         /// </summary>

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/QueryExecuteCompleteNotification.cs

@@ -21,11 +21,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
         /// Summaries of the result sets that were returned with the query
         /// </summary>
         public BatchSummary[] BatchSummaries { get; set; }
-
-        /// <summary>
-        /// Error message, if any
-        /// </summary>
-        public string Message { get; set; }
     }
 
     public class QueryExecuteCompleteEvent

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/QueryExecuteMessageNotification.cs

@@ -0,0 +1,32 @@
+//
+// Copyright (c) Microsoft. All rights reserved.
+// Licensed under the MIT license. See LICENSE file in the project root for full license information.
+//
+
+using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol.Contracts;
+
+namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
+{
+    /// <summary>
+    /// Parameters to be sent back with a message notification
+    /// </summary>
+    public class QueryExecuteMessageParams
+    {
+        /// <summary>
+        /// URI for the editor that owns the query
+        /// </summary>
+        public string OwnerUri { get; set; }
+
+        /// <summary>
+        /// The message that is being returned
+        /// </summary>
+        public ResultMessage Message { get; set; }
+    }
+
+    public class QueryExecuteMessageEvent
+    {
+        public static readonly
+            EventType<QueryExecuteMessageParams> Type =
+            EventType<QueryExecuteMessageParams>.Create("query/message");
+    }
+}

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/QueryExecuteRequest.cs

@@ -28,10 +28,6 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
     /// </summary>
     public class QueryExecuteResult
     {
-        /// <summary>
-        /// Informational messages from the query runner. Optional, can be set to null.
-        /// </summary>
-        public string Messages { get; set; }
     }
 
     public class QueryExecuteRequest

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Contracts/ResultMessage.cs

@@ -12,6 +12,17 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
     /// </summary>
     public class ResultMessage
     {
+        /// <summary>
+        /// ID of the batch that generated this message. If null, this message
+        /// was not generated as part of a batch
+        /// </summary>
+        public int? BatchId { get; set; }
+
+        /// <summary>
+        /// Whether or not this message is an error
+        /// </summary>
+        public bool IsError { get; set; }
+
         /// <summary>
         /// Timestamp of the message
         /// Stored in UTC ISO 8601 format; should be localized before displaying to any user
@@ -23,20 +34,13 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts
         /// </summary>
         public string Message { get; set; }
 
-        /// <summary>
-        /// Full constructor
-        /// </summary>
-        public ResultMessage(string timeStamp, string message)
-        {
-            Time = timeStamp;
-            Message = message;
-        }
-
         /// <summary>
         /// Constructor with default "Now" time
         /// </summary>
-        public ResultMessage(string message)
+        public ResultMessage(string message, bool isError, int? batchId)
         {
+            BatchId = batchId;
+            IsError = isError;
             Time = DateTime.Now.ToString("o");
             Message = message;
         }

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/Query.cs

@@ -99,6 +99,11 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
         /// </summary>
         public event Batch.BatchAsyncEventHandler BatchCompleted;
 
+        /// <summary>
+        /// Event that will be called when a message has been emitted
+        /// </summary>
+        public event Batch.BatchAsyncMessageHandler BatchMessageSent;
+
         /// <summary>
         /// Event to be called when a batch starts execution.
         /// </summary>
@@ -272,6 +277,14 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             // Don't actually execute if there aren't any batches to execute
             if (Batches.Length == 0)
             {
+                if (BatchMessageSent != null)
+                {
+                    await BatchMessageSent(new ResultMessage(SR.QueryServiceCompletedSuccessfully, false, null));
+                }
+                if (QueryCompleted != null)
+                {
+                    await QueryCompleted(this);
+                }
                 return;
             }
 
@@ -308,6 +321,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                     {
                         b.BatchStart += BatchStarted;
                         b.BatchCompletion += BatchCompleted;
+                        b.BatchMessageSent += BatchMessageSent;
                         b.ResultSetCompletion += ResultSetCompleted;
                         await b.Execute(conn, cancellationSource.Token);
                     }

src/Microsoft.SqlTools.ServiceLayer/QueryExecution/QueryExecutionService.cs

@@ -370,10 +370,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 }
 
                 // Send the result stating that the query was successfully started
-                await requestContext.SendResult(new QueryExecuteResult
+                await requestContext.SendResult(new QueryExecuteResult());
-                {
-                    Messages = newQuery.Batches.Length == 0 ? SR.QueryServiceCompletedSuccessfully : null
-                });
 
                 return newQuery;
             }
@@ -411,7 +408,7 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
                 QueryExecuteCompleteParams eventParams = new QueryExecuteCompleteParams
                 {
                     OwnerUri = executeParams.OwnerUri,
-                    Message = errorMessage              
+                    //Message = errorMessage              
                 };
                 await requestContext.SendEvent(QueryExecuteCompleteEvent.Type, eventParams);
             };
@@ -443,6 +440,17 @@ namespace Microsoft.SqlTools.ServiceLayer.QueryExecution
             };
             query.BatchCompleted += batchCompleteCallback;
 
+            Batch.BatchAsyncMessageHandler batchMessageCallback = async m =>
+            {
+                QueryExecuteMessageParams eventParams = new QueryExecuteMessageParams
+                {
+                    Message = m,
+                    OwnerUri = executeParams.OwnerUri
+                };
+                await requestContext.SendEvent(QueryExecuteMessageEvent.Type, eventParams);
+            };
+            query.BatchMessageSent += batchMessageCallback;
+
             // Setup the ResultSet completion callback
             ResultSet.ResultSetAsyncEventHandler resultCallback = async r =>
             {

test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/Connection/ReliableConnectionTests.cs

@@ -777,9 +777,6 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.Connection
                 var detectionStrategy2 = new TestSqlAzureTemporaryAndIgnorableErrorDetectionStrategy();
                 Assert.NotNull(detectionStrategy2.InvokeCanRetrySqlException(sqlException));
                 Assert.NotNull(detectionStrategy2.InvokeShouldIgnoreSqlException(sqlException));
-
-                Batch batch = new Batch(Common.StandardQuery, Common.SubsectionDocument, Common.Ordinal, Common.GetFileStreamFactory(null));
-                batch.UnwrapDbException(sqlException);
             }
 
             var unknownCodeReason = RetryPolicy.ThrottlingReason.FromReasonCode(-1);

test/Microsoft.SqlTools.ServiceLayer.PerfTests/Tests/QueryExecutionTests.cs

@@ -79,7 +79,7 @@ namespace Microsoft.SqlTools.ServiceLayer.PerfTests
                 };
 
                 var result = await testHelper.Driver.SendRequest(QueryExecuteRequest.Type, queryParams);
-                if (result != null && string.IsNullOrEmpty(result.Messages))
+                if (result != null)
                 {
                     TestTimer timer = new TestTimer() { PrintResult = true };
                     await Common.ExecuteWithTimeout(timer, 100000, async () => 

test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/Common.cs

@@ -97,6 +97,16 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution
             return output;
         }
 
+        public static Dictionary<string, string>[][] GetTestDataSet(int dataSets)
+        {
+            List<Dictionary<string, string>[]> output = new List<Dictionary<string, string>[]>();
+            for(int dataSet = 0; dataSet < dataSets; dataSet++)
+            {
+                output.Add(StandardTestData);
+            }
+            return output.ToArray();
+        }
+
         public static async Task AwaitExecution(QueryExecutionService service, QueryExecuteParams qeParams,
             RequestContext<QueryExecuteResult> requestContext)
         {

test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/Execution/BatchTests.cs

@@ -7,8 +7,7 @@ using System;
 using System.Collections.Generic;
 using System.Data;
 using System.Data.Common;
-using System.Data.SqlClient;
+using System.Diagnostics.CodeAnalysis;
-using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.SqlTools.ServiceLayer.Connection;
@@ -32,12 +31,10 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
 
             // ... It should not have executed and no error
             Assert.False(batch.HasExecuted, "The query should not have executed.");
-            Assert.False(batch.HasError, "The batch should not have an error");
 
             // ... The results should be empty
             Assert.Empty(batch.ResultSets);
             Assert.Empty(batch.ResultSummaries);
-            Assert.Empty(batch.ResultMessages);
 
             // ... The start line of the batch should be 0
             Assert.Equal(0, batch.Selection.StartLine);
@@ -46,318 +43,186 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             Assert.Equal(Common.Ordinal, batch.Id);
 
             // ... The summary should have the same info
-            Assert.False(batch.Summary.HasError);
             Assert.Equal(Common.Ordinal, batch.Summary.Id);
             Assert.Null(batch.Summary.ResultSetSummaries);
-            Assert.Null(batch.Summary.Messages);
             Assert.Equal(0, batch.Summary.Selection.StartLine);
             Assert.NotEqual(default(DateTime).ToString("o"), batch.Summary.ExecutionStart); // Should have been set at construction
             Assert.Null(batch.Summary.ExecutionEnd);
             Assert.Null(batch.Summary.ExecutionElapsed);
         }
 
-        /// <summary>
-        /// Note: This test also tests the start notification feature
-        /// </summary>
         [Fact]
-        public void BatchExecuteNoResultSets()
+        public async Task BatchExecuteNoResultSets()
         {
             // Setup: 
-            // ... Create a callback for batch start
+            // ... Keep track of callbacks being called
-            BatchSummary batchSummaryFromStart = null;
+            int batchStartCalls = 0;
-            Batch.BatchAsyncEventHandler batchStartCallback = b =>
+            int batchEndCalls = 0;
-            {
+            int resultSetCalls = 0;
-                batchSummaryFromStart = b.Summary;
+            List<ResultMessage> messages = new List<ResultMessage>();
-                return Task.FromResult(0);
-            };
-
-            // ... Create a callback for batch completion
-            BatchSummary batchSummaryFromCompletion = null;
-            Batch.BatchAsyncEventHandler batchCompleteCallback = b =>
-            {
-                batchSummaryFromCompletion = b.Summary;
-                return Task.FromResult(0);
-            };
-
-            // ... Create a callback for result completion
-            bool resultCallbackFired = false;
-            ResultSet.ResultSetAsyncEventHandler resultSetCallback = r =>
-            {
-                resultCallbackFired = true;
-                return Task.FromResult(0);
-            };
 
             // If I execute a query that should get no result sets
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
             Batch batch = new Batch(Common.StandardQuery, Common.SubsectionDocument, Common.Ordinal, fileStreamFactory);
-            batch.BatchStart += batchStartCallback;
+            BatchCallbackHelper(batch,
-            batch.BatchCompletion += batchCompleteCallback;
+                b => batchStartCalls++,
-            batch.ResultSetCompletion += resultSetCallback;
+                b => batchEndCalls++,
-            batch.Execute(GetConnection(Common.CreateTestConnectionInfo(null, false)), CancellationToken.None).Wait();
+                m => messages.Add(m),
+                r => resultSetCalls++);
+            await batch.Execute(GetConnection(Common.CreateTestConnectionInfo(null, false)), CancellationToken.None);
 
             // Then:
-            // ... It should have executed without error
+            // ... Callbacks should have been called the appropriate number of times
-            Assert.True(batch.HasExecuted, "The query should have been marked executed.");
+            Assert.Equal(1, batchStartCalls);
-            Assert.False(batch.HasError, "The batch should not have an error");
+            Assert.Equal(1, batchEndCalls);
+            Assert.Equal(0, resultSetCalls);
 
-            // ... The results should be empty
+            // ... The batch and the summary should be correctly assigned
-            Assert.Empty(batch.ResultSets);
+            ValidateBatch(batch, 0);
-            Assert.Empty(batch.ResultSummaries);
+            ValidateBatchSummary(batch);
-
+            ValidateMessages(batch, 1, messages);
-            // ... The results should not be null
-            Assert.NotNull(batch.ResultSets);
-            Assert.NotNull(batch.ResultSummaries);
-
-            // ... There should be a message for how many rows were affected
-            Assert.Equal(1, batch.ResultMessages.Count());
-
-            // ... The callback for batch start should have been called
-            // ... The info from it should have been basic
-            Assert.NotNull(batchSummaryFromStart);
-            Assert.False(batchSummaryFromStart.HasError);
-            Assert.Equal(Common.Ordinal, batchSummaryFromStart.Id);
-            Assert.Equal(Common.SubsectionDocument, batchSummaryFromStart.Selection);
-            Assert.True(DateTime.Parse(batchSummaryFromStart.ExecutionStart) > default(DateTime));
-            Assert.Null(batchSummaryFromStart.ResultSetSummaries);
-            Assert.Null(batchSummaryFromStart.Messages);
-            Assert.Null(batchSummaryFromStart.ExecutionElapsed);
-            Assert.Null(batchSummaryFromStart.ExecutionEnd);
-
-            // ... The callback for batch completion should have been fired
-            // ... The summary should match the expected info
-            Assert.NotNull(batchSummaryFromCompletion);
-            Assert.False(batchSummaryFromCompletion.HasError);
-            Assert.Equal(Common.Ordinal, batchSummaryFromCompletion.Id);
-            Assert.Equal(0, batchSummaryFromCompletion.ResultSetSummaries.Length);
-            Assert.Equal(1, batchSummaryFromCompletion.Messages.Length);
-            Assert.Equal(Common.SubsectionDocument, batchSummaryFromCompletion.Selection);
-            Assert.True(DateTime.Parse(batchSummaryFromCompletion.ExecutionStart) > default(DateTime));
-            Assert.True(DateTime.Parse(batchSummaryFromCompletion.ExecutionEnd) > default(DateTime));
-            Assert.NotNull(batchSummaryFromCompletion.ExecutionElapsed);
-
-            // ... The callback for the result set should NOT have been fired
-            Assert.False(resultCallbackFired);
         }
 
         [Fact]
-        public void BatchExecuteOneResultSet()
+        public async Task BatchExecuteOneResultSet()
         {
+            // Setup: 
+            // ... Keep track of callbacks being called
+            int batchStartCalls = 0;
+            int batchEndCalls = 0;
+            int resultSetCalls = 0;
+            List<ResultMessage> messages = new List<ResultMessage>();
+
+            // ... Build a data set to return
             const int resultSets = 1;
-            ConnectionInfo ci = Common.CreateTestConnectionInfo(new[] { Common.StandardTestData }, false);
+            ConnectionInfo ci = Common.CreateTestConnectionInfo(Common.GetTestDataSet(resultSets), false);
-
-            // Setup: Create a callback for batch completion
-            BatchSummary batchSummaryFromCallback = null;
-            Batch.BatchAsyncEventHandler batchCallback = b =>
-            {
-                batchSummaryFromCallback = b.Summary;
-                return Task.FromResult(0);
-            };
-
-            // ... Create a callback for result set completion
-            bool resultCallbackFired = false;
-            ResultSet.ResultSetAsyncEventHandler resultSetCallback = r =>
-            {
-                resultCallbackFired = true;
-                return Task.FromResult(0);
-            };
 
             // If I execute a query that should get one result set
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
             Batch batch = new Batch(Common.StandardQuery, Common.SubsectionDocument, Common.Ordinal, fileStreamFactory);
-            batch.BatchCompletion += batchCallback;
+            BatchCallbackHelper(batch,
-            batch.ResultSetCompletion += resultSetCallback;
+                b => batchStartCalls++,
-            batch.Execute(GetConnection(ci), CancellationToken.None).Wait();
+                b => batchEndCalls++,
+                m => messages.Add(m),
+                r => resultSetCalls++);
+            await batch.Execute(GetConnection(ci), CancellationToken.None);
 
             // Then:
-            // ... It should have executed without error
+            // ... Callbacks should have been called the appropriate number of times
-            Assert.True(batch.HasExecuted, "The batch should have been marked executed.");
+            Assert.Equal(1, batchStartCalls);
-            Assert.False(batch.HasError, "The batch should not have an error");
+            Assert.Equal(1, batchEndCalls);
+            Assert.Equal(1, resultSetCalls);
 
             // ... There should be exactly one result set
-            Assert.Equal(resultSets, batch.ResultSets.Count);
+            ValidateBatch(batch, resultSets);
-            Assert.Equal(resultSets, batch.ResultSummaries.Length);
+            ValidateBatchSummary(batch);
-
+            ValidateMessages(batch, 1, messages);
-            // ... Inside the result set should be with 5 rows
-            Assert.Equal(Common.StandardRows, batch.ResultSets.First().RowCount);
-            Assert.Equal(Common.StandardRows, batch.ResultSummaries[0].RowCount);
-
-            // ... Inside the result set should have 5 columns
-            Assert.Equal(Common.StandardColumns, batch.ResultSets.First().Columns.Length);
-            Assert.Equal(Common.StandardColumns, batch.ResultSummaries[0].ColumnInfo.Length);
-
-            // ... There should be a message for how many rows were affected
-            Assert.Equal(resultSets, batch.ResultMessages.Count());
-
-            // ... The callback for batch completion should have been fired
-            Assert.NotNull(batchSummaryFromCallback);
-
-            // ... The callback for resultset completion should have been fired
-            Assert.True(resultCallbackFired);  // We only want to validate that it happened, validation of the 
-                                               // summary is done in result set tests
         }
 
         [Fact]
-        public void BatchExecuteTwoResultSets()
+        public async Task BatchExecuteTwoResultSets()
         {
-            var dataset = new[] { Common.StandardTestData, Common.StandardTestData };
+            // Setup: 
-            int resultSets = dataset.Length;
+            // ... Keep track of callbacks being called
-            ConnectionInfo ci = Common.CreateTestConnectionInfo(dataset, false);
+            int batchStartCalls = 0;
+            int batchEndCalls = 0;
+            int resultSetCalls = 0;
+            List<ResultMessage> messages = new List<ResultMessage>();
 
-            // Setup: Create a callback for batch completion
+            // ... Build a data set to return
-            BatchSummary batchSummaryFromCallback = null;
+            const int resultSets = 2;
-            Batch.BatchAsyncEventHandler batchCallback = b =>
+            ConnectionInfo ci = Common.CreateTestConnectionInfo(Common.GetTestDataSet(resultSets), false);
-            {
-                batchSummaryFromCallback = b.Summary;
-                return Task.FromResult(0);
-            };
-
-            // ... Create a callback for resultset completion
-            int resultSummaryCount  = 0;
-            ResultSet.ResultSetAsyncEventHandler resultSetCallback = r =>
-            {
-                resultSummaryCount++;
-                return Task.FromResult(0);
-            };
 
             // If I execute a query that should get two result sets
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
             Batch batch = new Batch(Common.StandardQuery, Common.SubsectionDocument, Common.Ordinal, fileStreamFactory);
-            batch.BatchCompletion += batchCallback;
+            BatchCallbackHelper(batch,
-            batch.ResultSetCompletion += resultSetCallback;
+                b => batchStartCalls++,
-            batch.Execute(GetConnection(ci), CancellationToken.None).Wait();
+                b => batchEndCalls++,
+                m => messages.Add(m),
+                r => resultSetCalls++);
+            await batch.Execute(GetConnection(ci), CancellationToken.None);
 
             // Then:
+            // ... Callbacks should have been called the appropriate number of times
+            Assert.Equal(1, batchStartCalls);
+            Assert.Equal(1, batchEndCalls);
+            Assert.Equal(2, resultSetCalls);
+
             // ... It should have executed without error
-            Assert.True(batch.HasExecuted, "The batch should have been marked executed.");
+            ValidateBatch(batch, resultSets);
-            Assert.False(batch.HasError, "The batch should not have an error");
+            ValidateBatchSummary(batch);
-
+            ValidateMessages(batch, 1, messages);
-            // ... There should be exactly two result sets
-            Assert.Equal(resultSets, batch.ResultSets.Count());
-
-            foreach (ResultSet rs in batch.ResultSets)
-            {
-                // ... Each result set should have 5 rows
-                Assert.Equal(Common.StandardRows, rs.RowCount);
-
-                // ... Inside each result set should be 5 columns
-                Assert.Equal(Common.StandardColumns, rs.Columns.Length);
-            }
-
-            // ... There should be exactly two result set summaries
-            Assert.Equal(resultSets, batch.ResultSummaries.Length);
-
-            foreach (ResultSetSummary rs in batch.ResultSummaries)
-            {
-                // ... Inside each result summary, there should be 5 rows
-                Assert.Equal(Common.StandardRows, rs.RowCount);
-
-                // ... Inside each result summary, there should be 5 column definitions
-                Assert.Equal(Common.StandardColumns, rs.ColumnInfo.Length);
-            }
-
-            // ... The callback for batch completion should have been fired
-            Assert.NotNull(batchSummaryFromCallback);
-
-            // ... The callback for result set completion should have been fired
-            Assert.Equal(2, resultSummaryCount);
         }
 
         [Fact]
-        public void BatchExecuteInvalidQuery()
+        public async Task BatchExecuteInvalidQuery()
         {
-            // Setup:
+            // Setup: 
-            // ... Create a callback for batch start
+            // ... Keep track of callbacks being called
-            bool batchStartCalled = false;
+            int batchStartCalls = 0;
-            Batch.BatchAsyncEventHandler batchStartCallback = b =>
+            int batchEndCalls = 0;
-            {
+            List<ResultMessage> messages = new List<ResultMessage>();
-                batchStartCalled = true;
-                return Task.FromResult(0);
-            };
-             
-            // ... Create a callback for batch completion
-            BatchSummary batchSummaryFromCallback = null;
-            Batch.BatchAsyncEventHandler batchCompleteCallback = b =>
-            {
-                batchSummaryFromCallback = b.Summary;
-                return Task.FromResult(0);
-            };
-
-            // ... Create a callback that will fail the test if it's called
-            ResultSet.ResultSetAsyncEventHandler resultSetCallback = r =>
-            {
-                throw new Exception("ResultSet callback was called when it should not have been.");
-            };
-
-            ConnectionInfo ci = Common.CreateTestConnectionInfo(null, true);
 
             // If I execute a batch that is invalid
+            var ci = Common.CreateTestConnectionInfo(null, true);
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
             Batch batch = new Batch(Common.StandardQuery, Common.SubsectionDocument, Common.Ordinal, fileStreamFactory);
-            batch.BatchStart += batchStartCallback;
+            BatchCallbackHelper(batch,
-            batch.BatchCompletion += batchCompleteCallback;
+                b => batchStartCalls++,
-            batch.ResultSetCompletion += resultSetCallback;
+                b => batchEndCalls++,
-            batch.Execute(GetConnection(ci), CancellationToken.None).Wait();
+                m => messages.Add(m),
+                r => { throw new Exception("ResultSet callback was called when it should not have been."); });
+            await batch.Execute(GetConnection(ci), CancellationToken.None);
 
             // Then:
-            // ... It should have executed with error
+            // ... Callbacks should have been called the appropriate number of times
-            Assert.True(batch.HasExecuted);
+            Assert.Equal(1, batchStartCalls);
-            Assert.True(batch.HasError);
+            Assert.Equal(1, batchEndCalls);
 
-            // ... There should be no result sets
+            // ... It should have executed without error
-            Assert.Empty(batch.ResultSets);
+            ValidateBatch(batch, 0);
-            Assert.Empty(batch.ResultSummaries);
+            ValidateBatchSummary(batch);
 
-            // ... There should be plenty of messages for the error
+            // ... There should be one error message returned
-            Assert.NotEmpty(batch.ResultMessages);
+            Assert.Equal(1, messages.Count);
-
+            Assert.All(messages, m =>
-            // ... The callback for batch completion should have been fired
+            {
-            Assert.NotNull(batchSummaryFromCallback);
+                Assert.True(m.IsError);
-
+                Assert.Equal(batch.Id, m.BatchId);
-            // ... The callback for batch start should have been fired
+            });
-            Assert.True(batchStartCalled);
         }
 
         [Fact]
         public async Task BatchExecuteExecuted()
         {
-            ConnectionInfo ci = Common.CreateTestConnectionInfo(new[] { Common.StandardTestData }, false);
+            // Setup: Build a data set to return
+            const int resultSets = 1;
+            ConnectionInfo ci = Common.CreateTestConnectionInfo(Common.GetTestDataSet(resultSets), false);
 
             // If I execute a batch
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
             Batch batch = new Batch(Common.StandardQuery, Common.SubsectionDocument, Common.Ordinal, fileStreamFactory);
-            batch.Execute(GetConnection(ci), CancellationToken.None).Wait();
+            await batch.Execute(GetConnection(ci), CancellationToken.None);
 
             // Then:
             // ... It should have executed without error
             Assert.True(batch.HasExecuted, "The batch should have been marked executed.");
-            Assert.False(batch.HasError, "The batch should not have an error");
-
-            // Setup for part 2: 
-            // ... Create a callback for batch completion
-            Batch.BatchAsyncEventHandler completeCallback = b =>
-            {
-                throw new Exception("Batch completion callback should not have been called");
-            };
-
-            // ... Create a callback for batch start
-            Batch.BatchAsyncEventHandler startCallback = b =>
-            {
-                throw new Exception("Batch start callback should not have been called");
-            };
 
             // If I execute it again
             // Then:
             // ... It should throw an invalid operation exception
-            batch.BatchStart += startCallback;
+            BatchCallbackHelper(batch,
-            batch.BatchCompletion += completeCallback;
+                b => { throw new Exception("Batch start callback should not have been called"); },
-            await Assert.ThrowsAsync<InvalidOperationException>(() =>
+                b => { throw new Exception("Batch completion callback should not have been called"); },
-                batch.Execute(GetConnection(ci), CancellationToken.None));
+                m => { throw new Exception("Message callback should not have been called"); },
+                null);
+            await Assert.ThrowsAsync<InvalidOperationException>(
+                () => batch.Execute(GetConnection(ci), CancellationToken.None));
 
             // ... The data should still be available without error
-            Assert.False(batch.HasError, "The batch should not be in an error condition");
+            ValidateBatch(batch, resultSets);
-            Assert.True(batch.HasExecuted, "The batch should still be marked executed.");
+            ValidateBatchSummary(batch);
-            Assert.NotEmpty(batch.ResultSets);
-            Assert.NotEmpty(batch.ResultSummaries);
         }
 
         [Theory]
@@ -397,13 +262,20 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
         {
             // If:
             // ... I call the StatementCompletedHandler
+            Batch batch = new Batch(Common.StandardQuery, Common.SubsectionDocument, Common.Ordinal, Common.GetFileStreamFactory(null));
+            int messageCalls = 0;
+            batch.BatchMessageSent += args =>
+            {
+                messageCalls++;
+                return Task.FromResult(0);
+            };
+
             // Then:
-            // ... a ResultMaessage should be logged in the resultsMessages collection
+            // ... The message handler for the batch should havve been called twice
-             Batch batch = new Batch(Common.StandardQuery, Common.SubsectionDocument, Common.Ordinal, Common.GetFileStreamFactory(null));
+            batch.StatementCompletedHandler(null, new StatementCompletedEventArgs(1));
-             batch.StatementCompletedHandler(null, new StatementCompletedEventArgs(1));
+            Assert.True(messageCalls == 1);
-             Assert.True(batch.ResultMessages.Count() == 1);
+            batch.StatementCompletedHandler(null, new StatementCompletedEventArgs(2));
-             batch.StatementCompletedHandler(null, new StatementCompletedEventArgs(2));
+            Assert.True(messageCalls == 2);
-             Assert.True(batch.ResultMessages.Count() == 2);
         }     
       
         private static DbConnection GetConnection(ConnectionInfo info)
@@ -411,5 +283,81 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             return info.Factory.CreateSqlConnection(ConnectionService.BuildConnectionString(info.ConnectionDetails));
         }
 
+        [SuppressMessage("ReSharper", "UnusedParameter.Local")]
+        private static void ValidateBatch(Batch batch, int expectedResultSets)
+        {
+            // The batch should be executed
+            Assert.True(batch.HasExecuted, "The query should have been marked executed.");
+
+            // Result set list should never be null
+            Assert.NotNull(batch.ResultSets);
+            Assert.NotNull(batch.ResultSummaries);
+
+            // Make sure the number of result sets matches
+            Assert.Equal(expectedResultSets, batch.ResultSets.Count);
+            Assert.Equal(expectedResultSets, batch.ResultSummaries.Length);
+        }
+
+        private static void ValidateBatchSummary(Batch batch)
+        {
+            BatchSummary batchSummary = batch.Summary;
+
+            Assert.NotNull(batchSummary);
+            Assert.Equal(batch.Id, batchSummary.Id);
+            Assert.Equal(batch.ResultSets.Count, batchSummary.ResultSetSummaries.Length);
+            Assert.Equal(batch.Selection, batchSummary.Selection);
+
+            // Something other than default date is provided for start and end times
+            Assert.True(DateTime.Parse(batchSummary.ExecutionStart) > default(DateTime));   
+            Assert.True(DateTime.Parse(batchSummary.ExecutionEnd) > default(DateTime));
+            Assert.NotNull(batchSummary.ExecutionElapsed);
+        }
+
+        [SuppressMessage("ReSharper", "UnusedParameter.Local")]
+        private static void ValidateMessages(Batch batch, int expectedMessages, IList<ResultMessage> messages)
+        {
+            // There should be equal number of messages to result sets
+            Assert.Equal(expectedMessages, messages.Count);
+
+            // No messages should be errors
+            // All messages must have the batch ID
+            Assert.All(messages, m =>
+            {
+                Assert.False(m.IsError);
+                Assert.Equal(batch.Id, m.BatchId);
+            });
+        }
+
+        private static void BatchCallbackHelper(Batch batch, Action<Batch> startCallback, Action<Batch> endCallback,
+            Action<ResultMessage> messageCallback, Action<ResultSet> resultCallback)
+        {
+            // Setup the callback for batch start
+            batch.BatchStart += b =>
+            {
+                startCallback?.Invoke(b);
+                return Task.FromResult(0);
+            };
+
+            // Setup the callback for batch completion
+            batch.BatchCompletion += b =>
+            {
+                endCallback?.Invoke(b);
+                return Task.FromResult(0);
+            };
+
+            // Setup the callback for batch messages
+            batch.BatchMessageSent += (m) =>
+            {
+                messageCallback?.Invoke(m);
+                return Task.FromResult(0);
+            };
+
+            // Setup the result set completion callback
+            batch.ResultSetCompletion += r =>
+            {
+                resultCallback?.Invoke(r);
+                return Task.FromResult(0);
+            };
+        }
     }
 }

test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/Execution/QueryTests.cs

@@ -5,9 +5,12 @@
 
 using System;
 using System.Collections.Generic;
+using System.Linq;
 using System.Threading.Tasks;
+using Castle.Components.DictionaryAdapter;
 using Microsoft.SqlTools.ServiceLayer.Connection;
 using Microsoft.SqlTools.ServiceLayer.QueryExecution;
+using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
 using Microsoft.SqlTools.ServiceLayer.SqlContext;
 using Xunit;
 
@@ -16,6 +19,22 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
     public class QueryTests
     {
 
+        [Fact]
+        public void QueryCreationCorrect()
+        {
+            // If:
+            // ... I create a query
+            ConnectionInfo ci = Common.CreateTestConnectionInfo(null, false);
+            var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
+            Query query = new Query(Common.StandardQuery, ci, new QueryExecutionSettings(), fileStreamFactory);
+
+            // Then:
+            // ... I should get back two batches to execute that haven't been executed
+            Assert.NotEmpty(query.QueryText);
+            Assert.False(query.HasExecuted);
+            Assert.Throws<InvalidOperationException>(() => query.BatchSummaries);
+        }
+
         [Fact]
         public void QueryExecuteNoQueryText()
         {
@@ -63,44 +82,30 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
         public void QueryExecuteSingleBatch()
         {
             // Setup:
-            // ... Create a callback for atch start
+            // ... Keep track of how many times the callbacks were called
             int batchStartCallbacksReceived = 0;
-            Batch.BatchAsyncEventHandler batchStartCallback = b =>
-            {
-                batchStartCallbacksReceived++;
-                return Task.FromResult(0);
-            };
-
-            // ... Create a callback for batch completion
             int batchCompleteCallbacksReceived = 0;
-            Batch.BatchAsyncEventHandler batchCompleteCallback = summary =>
+            int batchMessageCallbacksReceived = 0;
-            {
-                batchCompleteCallbacksReceived++;
-                return Task.CompletedTask;
-            };
 
             // If:
             // ... I create a query from a single batch (without separator)
             ConnectionInfo ci = Common.CreateTestConnectionInfo(null, false);
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
             Query query = new Query(Common.StandardQuery, ci, new QueryExecutionSettings(), fileStreamFactory);
-            query.BatchStarted += batchStartCallback;
+            BatchCallbackHelper(query,
-            query.BatchCompleted += batchCompleteCallback;
+                b => batchStartCallbacksReceived++,
+                b => batchCompleteCallbacksReceived++,
+                m => batchMessageCallbacksReceived++);
 
-            // Then:
-            // ... I should get a single batch to execute that hasn't been executed
-            Assert.NotEmpty(query.QueryText);
-            Assert.NotEmpty(query.Batches);
-            Assert.Equal(1, query.Batches.Length);
-            Assert.False(query.HasExecuted);
-            Assert.Throws<InvalidOperationException>(() => query.BatchSummaries);
-
-            // If:
             // ... I then execute the query
             query.Execute();
             query.ExecutionTask.Wait();
 
             // Then:
+            // ... There should be exactly 1 batch
+            Assert.NotEmpty(query.Batches);
+            Assert.Equal(1, query.Batches.Length);
+
             // ... The query should have completed successfully with one batch summary returned
             Assert.True(query.HasExecuted);
             Assert.NotEmpty(query.BatchSummaries);
@@ -109,38 +114,24 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             // ... The batch callbacks should have been called precisely 1 time
             Assert.Equal(1, batchStartCallbacksReceived);
             Assert.Equal(1, batchCompleteCallbacksReceived);
+            Assert.Equal(1, batchMessageCallbacksReceived);
         }
 
         [Fact]
-        public void QueryExecuteNoOpBatch()
+        public void QueryExecuteSingleNoOpBatch()
         {
-            // Setup:
+            // Setup: Keep track of all the messages received
-            // ... Create a callback for batch startup
+            List<ResultMessage> messages = new List<ResultMessage>();
-            Batch.BatchAsyncEventHandler batchStartCallback = b =>
-            {
-                throw new Exception("Batch startup callback should not have been called.");
-            };
-
-            // ... Create a callback for batch completion
-            Batch.BatchAsyncEventHandler batchCompletionCallback = summary =>
-            {
-                throw new Exception("Batch completion callback was called");
-            };
 
             // If:
             // ... I create a query from a single batch that does nothing
             ConnectionInfo ci = Common.CreateTestConnectionInfo(null, false);
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
             Query query = new Query(Common.NoOpQuery, ci, new QueryExecutionSettings(), fileStreamFactory);
-            query.BatchStarted += batchStartCallback;
+            BatchCallbackHelper(query,
-            query.BatchCompleted += batchCompletionCallback;
+                b => { throw new Exception("Batch startup callback should not have been called."); },
-
+                b => { throw new Exception("Batch completion callback was called"); },
-            // Then:
+                m => messages.Add(m));
-            // ... I should get no batches back
-            Assert.NotEmpty(query.QueryText);
-            Assert.Empty(query.Batches);
-            Assert.False(query.HasExecuted);
-            Assert.Throws<InvalidOperationException>(() => query.BatchSummaries);
 
             // If:
             // ... I Then execute the query
@@ -148,30 +139,27 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             query.ExecutionTask.Wait();
 
             // Then:
+            // ... There should be no batches
+            Assert.Empty(query.Batches);
+
             // ... The query should have completed successfully with no batch summaries returned
             Assert.True(query.HasExecuted);
             Assert.Empty(query.BatchSummaries);
+
+            // ... The message callback should have been called exactly once
+            // ... The message must not have a batch associated with it
+            Assert.Equal(1, messages.Count);
+            Assert.Null(messages[0].BatchId);
         }
 
         [Fact]
-        public void QueryExecuteMultipleBatches()
+        public void QueryExecuteMultipleResultBatches()
         {
             // Setup:
-            // ... Create a callback for batch start
+            // ... Keep track of how many callbacks are received
             int batchStartCallbacksReceived = 0;
-            Batch.BatchAsyncEventHandler batchStartCallback = b =>
-            {
-                batchStartCallbacksReceived++;
-                return Task.FromResult(0);
-            };
-
-            // ... Create a callback for batch completion
             int batchCompletedCallbacksReceived = 0;
-            Batch.BatchAsyncEventHandler batchCompletedCallback = summary =>
+            int batchMessageCallbacksReceived = 0;
-            {
-                batchCompletedCallbacksReceived++;
-                return Task.FromResult(0);
-            };
 
             // If:
             // ... I create a query from two batches (with separator)
@@ -179,52 +167,39 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             string queryText = string.Format("{0}\r\nGO\r\n{0}", Common.StandardQuery);
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
             Query query = new Query(queryText, ci, new QueryExecutionSettings(), fileStreamFactory);
-            query.BatchStarted += batchStartCallback;
+            BatchCallbackHelper(query,
-            query.BatchCompleted += batchCompletedCallback;
+                b => batchStartCallbacksReceived++,
+                b => batchCompletedCallbacksReceived++,
+                m => batchMessageCallbacksReceived++);
 
-            // Then:
-            // ... I should get back two batches to execute that haven't been executed
-            Assert.NotEmpty(query.QueryText);
-            Assert.NotEmpty(query.Batches);
-            Assert.Equal(2, query.Batches.Length);
-            Assert.False(query.HasExecuted);
-            Assert.Throws<InvalidOperationException>(() => query.BatchSummaries);
-
-            // If:
             // ... I then execute the query
             query.Execute();
             query.ExecutionTask.Wait();
 
             // Then:
+            // ... I should get back a query with one batch (no op batch is not included)
+            Assert.NotEmpty(query.Batches);
+            Assert.Equal(2, query.Batches.Length);
+
             // ... The query should have completed successfully with two batch summaries returned
             Assert.True(query.HasExecuted);
             Assert.NotEmpty(query.BatchSummaries);
             Assert.Equal(2, query.BatchSummaries.Length);
 
-            // ... The batch start and completion callbacks should have been called precisely 2 times
+            // ... The batch start, complete, and message callbacks should have been called precisely 2 times
             Assert.Equal(2, batchStartCallbacksReceived);
             Assert.Equal(2, batchCompletedCallbacksReceived);
+            Assert.Equal(2, batchMessageCallbacksReceived);
         }
 
         [Fact]
         public void QueryExecuteMultipleBatchesWithNoOp()
         {
             // Setup:
-            // ... Create a callback for batch start
+            // ... Keep track of how many times callbacks are called
             int batchStartCallbacksReceived = 0;
-            Batch.BatchAsyncEventHandler batchStartCallback = b =>
-            {
-                batchStartCallbacksReceived++;
-                return Task.FromResult(0);
-            };
-
-            // ... Create a callback for batch completion
             int batchCompletionCallbacksReceived = 0;
-            Batch.BatchAsyncEventHandler batchCompletionCallback = summary =>
+            int batchMessageCallbacksReceived = 0;
-            {
-                batchCompletionCallbacksReceived++;
-                return Task.CompletedTask;
-            };
 
             // If:
             // ... I create a query from a two batches (with separator)
@@ -232,22 +207,20 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             string queryText = string.Format("{0}\r\nGO\r\n{1}", Common.StandardQuery, Common.NoOpQuery);
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
             Query query = new Query(queryText, ci, new QueryExecutionSettings(), fileStreamFactory);
-            query.BatchStarted += batchStartCallback;
+            BatchCallbackHelper(query, 
-            query.BatchCompleted += batchCompletionCallback;
+                b => batchStartCallbacksReceived++,
+                b => batchCompletionCallbacksReceived++,
+                m => batchMessageCallbacksReceived++);
 
-            // Then:
-            // ... I should get back one batch to execute that hasn't been executed
-            Assert.NotEmpty(query.QueryText);
-            Assert.NotEmpty(query.Batches);
-            Assert.Equal(1, query.Batches.Length);
-            Assert.False(query.HasExecuted);
-            Assert.Throws<InvalidOperationException>(() => query.BatchSummaries);
-
-            // If:
             // .. I then execute the query
             query.Execute();
             query.ExecutionTask.Wait();
 
+            // Then:
+            // ... I should get back a query with one batch (no op batch is not included)
+            Assert.NotEmpty(query.Batches);
+            Assert.Equal(1, query.Batches.Length);
+
             // ... The query should have completed successfully with one batch summary returned
             Assert.True(query.HasExecuted);
             Assert.NotEmpty(query.BatchSummaries);
@@ -256,61 +229,107 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             // ... The batch callbacks should have been called precisely 1 time
             Assert.Equal(1, batchStartCallbacksReceived);
             Assert.Equal(1, batchCompletionCallbacksReceived);
+            Assert.Equal(1, batchMessageCallbacksReceived);
+        }
+
+        [Fact]
+        public async Task QueryExecuteMultipleNoOpBatches()
+        {
+            // Setup:
+            // ... Keep track of how many messages were sent
+            List<ResultMessage> messages = new List<ResultMessage>();
+
+            // If:
+            // ... I create a query from a two batches (with separator)
+            ConnectionInfo ci = Common.CreateTestConnectionInfo(null, false);
+            string queryText = string.Format("{0}\r\nGO\r\n{1}", Common.NoOpQuery, Common.NoOpQuery);
+            var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
+            Query query = new Query(queryText, ci, new QueryExecutionSettings(), fileStreamFactory);
+            BatchCallbackHelper(query,
+                b => { throw new Exception("Batch start handler was called"); },
+                b => { throw new Exception("Batch completed handler was called"); },
+                m => messages.Add(m));
+
+            // .. I then execute the query
+            query.Execute();
+            await query.ExecutionTask;
+
+            // Then:
+            // ... I should get back a query with no batches
+            Assert.Empty(query.Batches);
+
+            // ... The query should have completed successfully with one zero batch summaries returned
+            Assert.True(query.HasExecuted);
+            Assert.Empty(query.BatchSummaries);
+
+            // ... The message callback should have been called exactly once
+            // ... The message must not have a batch associated with it
+            Assert.Equal(1, messages.Count);
+            Assert.Null(messages[0].BatchId);
         }
 
         [Fact]
         public void QueryExecuteInvalidBatch()
         {
             // Setup:
-            // ... Create a callback for batch start
+            // ... Keep track of how many times a method is called
             int batchStartCallbacksReceived = 0;
-            Batch.BatchAsyncEventHandler batchStartCallback = b =>
-            {
-                batchStartCallbacksReceived++;
-                return Task.FromResult(0);
-            };
-
-            // ... Create a callback for batch completion
             int batchCompletionCallbacksReceived = 0;
-            Batch.BatchAsyncEventHandler batchCompltionCallback = summary =>
+            List<ResultMessage> messages = new List<ResultMessage>();
-            {
-                batchCompletionCallbacksReceived++;
-                return Task.CompletedTask;
-            };
 
             // If:
             // ... I create a query from an invalid batch
             ConnectionInfo ci = Common.CreateTestConnectionInfo(null, true);
             var fileStreamFactory = Common.GetFileStreamFactory(new Dictionary<string, byte[]>());
             Query query = new Query(Common.InvalidQuery, ci, new QueryExecutionSettings(), fileStreamFactory);
-            query.BatchStarted += batchStartCallback;
+            BatchCallbackHelper(query,
-            query.BatchCompleted += batchCompltionCallback;
+                b => batchStartCallbacksReceived++,
+                b => batchCompletionCallbacksReceived++,
+                m => messages.Add(m));
 
-            // Then:
-            // ... I should get back a query with one batch not executed
-            Assert.NotEmpty(query.QueryText);
-            Assert.NotEmpty(query.Batches);
-            Assert.Equal(1, query.Batches.Length);
-            Assert.False(query.HasExecuted);
-            Assert.Throws<InvalidOperationException>(() => query.BatchSummaries);
-
-            // If:
             // ... I then execute the query
             query.Execute();
             query.ExecutionTask.Wait();
 
             // Then:
+            // ... I should get back a query with one batch
+            Assert.NotEmpty(query.Batches);
+            Assert.Equal(1, query.Batches.Length);
+
             // ... There should be an error on the batch
             Assert.True(query.HasExecuted);
             Assert.NotEmpty(query.BatchSummaries);
             Assert.Equal(1, query.BatchSummaries.Length);
-            Assert.True(query.BatchSummaries[0].HasError);
+            Assert.True(messages.Any(m => m.IsError));
-            Assert.NotEmpty(query.BatchSummaries[0].Messages);
 
             // ... The batch callbacks should have been called once
             Assert.Equal(1, batchStartCallbacksReceived);
             Assert.Equal(1, batchCompletionCallbacksReceived);
         }
 
+        private static void BatchCallbackHelper(Query q, Action<Batch> startCallback, Action<Batch> endCallback,
+            Action<ResultMessage> messageCallback)
+        {
+            // Setup the callback for batch start
+            q.BatchStarted += b =>
+            {
+                startCallback?.Invoke(b);
+                return Task.FromResult(0);
+            };
+
+            // Setup the callback for batch completion
+            q.BatchCompleted += b =>
+            {
+                endCallback?.Invoke(b);
+                return Task.FromResult(0);
+            };
+
+            // Setup the callback for batch messages
+            q.BatchMessageSent += (m) =>
+            {
+                messageCallback?.Invoke(m);
+                return Task.FromResult(0);
+            };
+        }
     }
 }

test/Microsoft.SqlTools.ServiceLayer.Test/QueryExecution/Execution/ServiceIntegrationTests.cs

@@ -21,14 +21,18 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             // ... I request to execute a valid query with all batches as no op
             var workspaceService = GetDefaultWorkspaceService(string.Format("{0}\r\nGO\r\n{0}", Common.NoOpQuery));
             var queryService = Common.GetPrimedExecutionService(null, true, false, workspaceService);
-            var queryParams = new QueryExecuteParams { QuerySelection = Common.WholeDocument, OwnerUri = Common.OwnerUri };
+            var queryParams = new QueryExecuteParams {QuerySelection = Common.WholeDocument, OwnerUri = Common.OwnerUri};
 
             var efv = new EventFlowValidator<QueryExecuteResult>()
-                .AddResultValidation(p =>
+                .AddStandardQueryResultValidator()
+                .AddStandardMessageValidator()
+                .AddEventValidation(QueryExecuteCompleteEvent.Type, p =>
                 {
-                    Assert.False(string.IsNullOrWhiteSpace(p.Messages));
+                    // Validate OwnerURI matches
-                })
+                    Assert.Equal(Common.OwnerUri, p.OwnerUri);
-                .Complete();
+                    Assert.NotNull(p.BatchSummaries);
+                    Assert.Equal(0, p.BatchSummaries.Length);
+                }).Complete();
             await Common.AwaitExecution(queryService, queryParams, efv.Object);
 
             // Then:
@@ -38,7 +42,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             // ... There should be one active query
             Assert.Equal(1, queryService.ActiveQueries.Count);
         }
-    
+
         [Fact]
         public async Task QueryExecuteSingleBatchNoResultsTest()
         {
@@ -46,11 +50,12 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             // ... I request to execute a valid query with no results
             var workspaceService = GetDefaultWorkspaceService(Common.StandardQuery);
             var queryService = Common.GetPrimedExecutionService(null, true, false, workspaceService);
-            var queryParams = new QueryExecuteParams { QuerySelection = Common.WholeDocument, OwnerUri = Common.OwnerUri };
+            var queryParams = new QueryExecuteParams {QuerySelection = Common.WholeDocument, OwnerUri = Common.OwnerUri};
 
             var efv = new EventFlowValidator<QueryExecuteResult>()
                 .AddStandardQueryResultValidator()
                 .AddStandardBatchStartValidator()
+                .AddStandardMessageValidator()
                 .AddStandardBatchCompleteValidator()
                 .AddStandardQueryCompleteValidator(1)
                 .Complete();
@@ -71,18 +76,20 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             // If:
             // ... I request to execute a valid query with results
             var workspaceService = GetDefaultWorkspaceService(Common.StandardQuery);
-            var queryService = Common.GetPrimedExecutionService(new[] { Common.StandardTestData }, true, false, workspaceService);
+            var queryService = Common.GetPrimedExecutionService(new[] {Common.StandardTestData}, true, false,
-            var queryParams = new QueryExecuteParams { OwnerUri = Common.OwnerUri, QuerySelection = Common.WholeDocument };
+                workspaceService);
+            var queryParams = new QueryExecuteParams {OwnerUri = Common.OwnerUri, QuerySelection = Common.WholeDocument};
 
             var efv = new EventFlowValidator<QueryExecuteResult>()
                 .AddStandardQueryResultValidator()
                 .AddStandardBatchStartValidator()
                 .AddStandardResultSetValidator()
+                .AddStandardMessageValidator()
                 .AddStandardBatchCompleteValidator()
                 .AddStandardQueryCompleteValidator(1)
                 .Complete();
             await Common.AwaitExecution(queryService, queryParams, efv.Object);
-            
+
             // Then:
             // ... All events should have been called as per their flow validator
             efv.Validate();
@@ -97,15 +104,16 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             // If:
             // ... I request to execute a valid query with one batch and multiple result sets
             var workspaceService = GetDefaultWorkspaceService(Common.StandardQuery);
-            var dataset = new[] { Common.StandardTestData, Common.StandardTestData };
+            var dataset = new[] {Common.StandardTestData, Common.StandardTestData};
             var queryService = Common.GetPrimedExecutionService(dataset, true, false, workspaceService);
-            var queryParams = new QueryExecuteParams { OwnerUri = Common.OwnerUri, QuerySelection = Common.WholeDocument };
+            var queryParams = new QueryExecuteParams {OwnerUri = Common.OwnerUri, QuerySelection = Common.WholeDocument};
 
             var efv = new EventFlowValidator<QueryExecuteResult>()
                 .AddStandardQueryResultValidator()
                 .AddStandardBatchStartValidator()
                 .AddStandardResultSetValidator()
                 .AddStandardResultSetValidator()
+                .AddStandardMessageValidator()
                 .AddStandardQueryCompleteValidator(1)
                 .Complete();
             await Common.AwaitExecution(queryService, queryParams, efv.Object);
@@ -124,17 +132,19 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             // If:
             // ... I request a to execute a valid query with multiple batches
             var workspaceService = GetDefaultWorkspaceService(string.Format("{0}\r\nGO\r\n{0}", Common.StandardQuery));
-            var dataSet = new[] { Common.StandardTestData };
+            var dataSet = new[] {Common.StandardTestData};
             var queryService = Common.GetPrimedExecutionService(dataSet, true, false, workspaceService);
-            var queryParams = new QueryExecuteParams { OwnerUri = Common.OwnerUri, QuerySelection = Common.WholeDocument };
+            var queryParams = new QueryExecuteParams {OwnerUri = Common.OwnerUri, QuerySelection = Common.WholeDocument};
 
             var efv = new EventFlowValidator<QueryExecuteResult>()
                 .AddStandardQueryResultValidator()
                 .AddStandardBatchStartValidator()
                 .AddStandardResultSetValidator()
+                .AddStandardMessageValidator()
                 .AddStandardBatchCompleteValidator()
                 .AddStandardBatchCompleteValidator()
                 .AddStandardResultSetValidator()
+                .AddStandardMessageValidator()
                 .AddStandardBatchCompleteValidator()
                 .AddStandardQueryCompleteValidator(2)
                 .Complete();
@@ -156,7 +166,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             // ... I request to execute a query using a file URI that isn't connected
             var workspaceService = GetDefaultWorkspaceService(Common.StandardQuery);
             var queryService = Common.GetPrimedExecutionService(null, false, false, workspaceService);
-            var queryParams = new QueryExecuteParams { OwnerUri = "notConnected", QuerySelection = Common.WholeDocument };
+            var queryParams = new QueryExecuteParams {OwnerUri = "notConnected", QuerySelection = Common.WholeDocument};
 
             var efv = new EventFlowValidator<QueryExecuteResult>()
                 .AddErrorValidation<string>(Assert.NotEmpty)
@@ -178,14 +188,14 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             // ... I request to execute a query
             var workspaceService = GetDefaultWorkspaceService(Common.StandardQuery);
             var queryService = Common.GetPrimedExecutionService(null, true, false, workspaceService);
-            var queryParams = new QueryExecuteParams { OwnerUri = Common.OwnerUri, QuerySelection = Common.WholeDocument };
+            var queryParams = new QueryExecuteParams {OwnerUri = Common.OwnerUri, QuerySelection = Common.WholeDocument};
 
             // Note, we don't care about the results of the first request
             var firstRequestContext = RequestContextMocks.Create<QueryExecuteResult>(null);
             await Common.AwaitExecution(queryService, queryParams, firstRequestContext.Object);
 
             // ... And then I request another query without waiting for the first to complete
-            queryService.ActiveQueries[Common.OwnerUri].HasExecuted = false;   // Simulate query hasn't finished
+            queryService.ActiveQueries[Common.OwnerUri].HasExecuted = false; // Simulate query hasn't finished
             var efv = new EventFlowValidator<QueryExecuteResult>()
                 .AddErrorValidation<string>(Assert.NotEmpty)
                 .Complete();
@@ -206,7 +216,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             // ... I request to execute a query
             var workspaceService = GetDefaultWorkspaceService(Common.StandardQuery);
             var queryService = Common.GetPrimedExecutionService(null, true, false, workspaceService);
-            var queryParams = new QueryExecuteParams { OwnerUri = Common.OwnerUri, QuerySelection = Common.WholeDocument };
+            var queryParams = new QueryExecuteParams {OwnerUri = Common.OwnerUri, QuerySelection = Common.WholeDocument};
 
             // Note, we don't care about the results of the first request
             var firstRequestContext = RequestContextMocks.Create<QueryExecuteResult>(null);
@@ -240,7 +250,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             // If:
             // ... I request to execute a query with a missing query string
             var queryService = Common.GetPrimedExecutionService(null, true, false, workspaceService);
-            var queryParams = new QueryExecuteParams { OwnerUri = Common.OwnerUri, QuerySelection = null };
+            var queryParams = new QueryExecuteParams {OwnerUri = Common.OwnerUri, QuerySelection = null};
 
             var efv = new EventFlowValidator<QueryExecuteResult>()
                 .AddErrorValidation<string>(Assert.NotEmpty)
@@ -262,7 +272,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             // ... I request to execute a query that is invalid
             var workspaceService = GetDefaultWorkspaceService(Common.StandardQuery);
             var queryService = Common.GetPrimedExecutionService(null, true, true, workspaceService);
-            var queryParams = new QueryExecuteParams { OwnerUri = Common.OwnerUri, QuerySelection = Common.WholeDocument };
+            var queryParams = new QueryExecuteParams {OwnerUri = Common.OwnerUri, QuerySelection = Common.WholeDocument};
 
             var efv = new EventFlowValidator<QueryExecuteResult>()
                 .AddStandardQueryResultValidator()
@@ -270,7 +280,6 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
                 .AddStandardBatchCompleteValidator()
                 .AddStandardQueryCompleteValidator(1)
                 .Complete();
-
             await Common.AwaitExecution(queryService, queryParams, efv.Object);
 
             // Then:
@@ -295,10 +304,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
             this EventFlowValidator<QueryExecuteResult> efv)
         {
             // We just need to makes sure we get a result back, there's no params to validate
-            return efv.AddResultValidation(r =>
+            return efv.AddResultValidation(Assert.NotNull);
-            {
-                Assert.Null(r.Messages);
-            });
         }
 
         public static EventFlowValidator<TRequestContext> AddStandardBatchStartValidator<TRequestContext>(
@@ -317,18 +323,29 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
         {
             return efv.AddEventValidation(QueryExecuteBatchCompleteEvent.Type, p =>
             {
-                // Validate OwnerURI and batch summary are returned
+                // Validate OwnerURI and result summary are returned
                 Assert.Equal(Common.OwnerUri, p.OwnerUri);
                 Assert.NotNull(p.BatchSummary);
             });
         }
 
+        public static EventFlowValidator<TRequestContext> AddStandardMessageValidator<TRequestContext>(
+            this EventFlowValidator<TRequestContext> efv)
+        {
+            return efv.AddEventValidation(QueryExecuteMessageEvent.Type, p =>
+            {
+                // Validate OwnerURI and message are returned
+                Assert.Equal(Common.OwnerUri, p.OwnerUri);
+                Assert.NotNull(p.Message);
+            });
+        }
+
         public static EventFlowValidator<TRequestContext> AddStandardResultSetValidator<TRequestContext>(
             this EventFlowValidator<TRequestContext> efv)
         {
             return efv.AddEventValidation(QueryExecuteResultSetCompleteEvent.Type, p =>
             {
-                // Validate OwnerURI and result summary are returned
+                // Validate OwnerURI and summary are returned
                 Assert.Equal(Common.OwnerUri, p.OwnerUri);
                 Assert.NotNull(p.ResultSetSummary);
             });
@@ -339,7 +356,6 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.QueryExecution.Execution
         {
             return efv.AddEventValidation(QueryExecuteCompleteEvent.Type, p =>
             {
-                Assert.True(string.IsNullOrWhiteSpace(p.Message));
                 Assert.Equal(Common.OwnerUri, p.OwnerUri);
                 Assert.NotNull(p.BatchSummaries);
                 Assert.Equal(expectedBatches, p.BatchSummaries.Length);

test/Microsoft.SqlTools.ServiceLayer.TestDriver/Tests/QueryExecutionTests.cs

@@ -327,34 +327,33 @@ namespace Microsoft.SqlTools.ServiceLayer.TestDriver.Tests
         }
         */
 
-        [Fact]
+        [Theory]
-        public async Task NoOpQueryReturnsMessage()
+        [InlineData("-- no-op")]
+        [InlineData("GO")]
+        [InlineData("GO -- no-op")]
+        public async Task NoOpQueryReturnsMessage(string query)
         {
-            // Given queries that do nothing (no-ops)...
-            var queries = new string[]
-            {
-                "-- no-op",
-                "GO",
-                "GO -- no-op"
-            };
-
             using (SelfCleaningTempFile queryTempFile = new SelfCleaningTempFile())
             using (TestHelper testHelper = new TestHelper())
             {
-                foreach (var query in queries)
+                Assert.True(await testHelper.Connect(queryTempFile.FilePath, ConnectionTestUtils.LocalhostConnection));
-                {
-                    Assert.True(await testHelper.Connect(queryTempFile.FilePath, ConnectionTestUtils.LocalhostConnection));
 
-                    // If the queries are executed...
+                // If: the query is executed...
-                    var queryResult = await testHelper.RunQueryAsync(queryTempFile.FilePath, query);
+                var queryResult = await testHelper.RunQueryAsync(queryTempFile.FilePath, query);
+                var message = await testHelper.WaitForMessage();
 
-                    // Then I expect messages that the commands were completed successfully to be in the result
+                // Then:
-                    Assert.NotNull(queryResult);
+                // ... I expect a query result to indicate successfully started query
-                    Assert.NotNull(queryResult.Messages);
+                Assert.NotNull(queryResult);
-                    Assert.Equal("Commands completed successfully.", queryResult.Messages);
 
-                    await testHelper.Disconnect(queryTempFile.FilePath);
+                // ... I expect a non-error message to be returned without a batch associated with it
-                }
+                Assert.NotNull(message);
+                Assert.NotNull(message.Message);
+                Assert.NotNull(message.Message.Message);
+                Assert.False(message.Message.IsError);
+                Assert.Null(message.Message.BatchId);
+
+                await testHelper.Disconnect(queryTempFile.FilePath);
             }
         }
     }

test/Microsoft.SqlTools.ServiceLayer.TestDriver/Tests/TestHelper.cs

@@ -295,7 +295,7 @@ namespace Microsoft.SqlTools.ServiceLayer.TestDriver.Tests
             };
 
             var result = await Driver.SendRequest(QueryExecuteRequest.Type, queryParams);
-            if (result != null && string.IsNullOrEmpty(result.Messages))
+            if (result != null)
             {
                 var eventResult = await Driver.WaitForEvent(QueryExecuteCompleteEvent.Type, timeoutMilliseconds);
                 return eventResult;
@@ -384,6 +384,15 @@ namespace Microsoft.SqlTools.ServiceLayer.TestDriver.Tests
             return result;
         }
 
+        /// <summary>
+        /// Waits for a message to be returned by the service
+        /// </summary>
+        /// <returns>A message from the service layer</returns>
+        public async Task<QueryExecuteMessageParams> WaitForMessage()
+        {
+            return await Driver.WaitForEvent(QueryExecuteMessageEvent.Type);
+        }
+
         public void WriteToFile(string ownerUri, string query)
         {
             lock (fileLock)