Added first batch of query execution tests using the test driver (#131)

* Added first batch of query execution tests using the test driver

* Fix issues from merge with dev

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

@@ -3,9 +3,7 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 //
 
-using System;
 using System.Collections.Generic;
-using System.Diagnostics;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
@@ -18,6 +16,190 @@ namespace Microsoft.SqlTools.ServiceLayer.TestDriver.Tests
 {
     public class QueryExecutionTests : TestBase
     {
+       [Fact]
+        public async Task TestQueryCancelReliability()
+        {
+            string ownerUri = System.IO.Path.GetTempFileName();
+            string query = "SELECT * FROM sys.objects a CROSS JOIN sys.objects b CROSS JOIN sys.objects c";
+
+            await Connect(ownerUri, ConnectionTestUtils.AzureTestServerConnection);
+
+            // Run and cancel 100 queries
+            for (int i = 0; i < 100; i++) 
+            {
+                var queryTask = RunQuery(ownerUri, query);
+
+                var cancelResult = await CancelQuery(ownerUri);
+                Assert.NotNull(cancelResult);
+                Assert.True(string.IsNullOrEmpty(cancelResult.Messages));
+
+                await queryTask;
+            }
+            
+            await Disconnect(ownerUri);
+        }
+
+        [Fact]
+        public async Task TestQueryDoesNotBlockOtherRequests()
+        {
+            string ownerUri = System.IO.Path.GetTempFileName();
+            string query = "SELECT * FROM sys.objects a CROSS JOIN sys.objects b CROSS JOIN sys.objects c";
+
+            await Connect(ownerUri, ConnectionTestUtils.AzureTestServerConnection);
+
+            // Start a long-running query
+            var queryTask = RunQuery(ownerUri, query, 60000);
+
+            // Interact with the service. None of these requests should time out while waiting for the query to finish
+            for (int i = 0; i < 10; i++)
+            {
+                string ownerUri2 = System.IO.Path.GetTempFileName();
+
+                await Connect(ownerUri2, ConnectionTestUtils.AzureTestServerConnection);
+                Assert.NotNull(await RequestCompletion(ownerUri2, "SELECT * FROM sys.objects", 0, 10));
+                await Disconnect(ownerUri2);
+            }
+
+            await CancelQuery(ownerUri);
+            await Disconnect(ownerUri);
+        }
+
+        [Fact]
+        public async Task TestParallelQueryExecution()
+        {
+            int queryCount = 10;
+
+            // Create n connections
+            string[] ownerUris = new string[queryCount];
+            for (int i = 0; i < queryCount; i++)
+            {
+                ownerUris[i] = System.IO.Path.GetTempFileName();
+                Assert.NotNull(await Connect(ownerUris[i], ConnectionTestUtils.AzureTestServerConnection));
+            }
+
+            // Run n queries at once
+            string query = "SELECT * FROM sys.objects";
+            var queryTasks = new Task<QueryExecuteCompleteParams>[queryCount];
+            for (int i = 0; i < queryCount; i++)
+            {
+                queryTasks[i] = RunQuery(ownerUris[i], query);
+            }
+            await Task.WhenAll(queryTasks);
+
+            // Verify that they all completed with results and Disconnect
+            for (int i = 0; i < queryCount; i++)
+            {
+                Assert.NotNull(queryTasks[i].Result);
+                Assert.NotNull(queryTasks[i].Result.BatchSummaries);
+                await Disconnect(ownerUris[i]);
+            }
+        }
+
+        [Fact]
+        public async Task TestSaveResultsDoesNotBlockOtherRequests()
+        {
+            string ownerUri = System.IO.Path.GetTempFileName();
+            string query = "SELECT * FROM sys.objects";
+
+            await Connect(ownerUri, ConnectionTestUtils.AzureTestServerConnection);
+
+            // Execute a query
+            await RunQuery(ownerUri, query);
+
+            // Spawn several tasks to save results
+            var saveTasks = new Task<SaveResultRequestResult>[100];
+            for (int i = 0; i < 100; i++)
+            {
+                if (i % 2 == 0)
+                {
+                    saveTasks[i] = SaveAsCsv(ownerUri, System.IO.Path.GetTempFileName(), 0, 0);
+                }
+                else
+                {
+                    saveTasks[i] = SaveAsJson(ownerUri, System.IO.Path.GetTempFileName(), 0, 0);
+                }
+            }
+
+            // Interact with the service. None of these requests should time out while waiting for the save results tasks to finish
+            for (int i = 0; i < 10; i++)
+            {
+                string ownerUri2 = System.IO.Path.GetTempFileName();
+
+                await Connect(ownerUri2, ConnectionTestUtils.AzureTestServerConnection);
+                Assert.NotNull(await RequestCompletion(ownerUri2, "SELECT * FROM sys.objects", 0, 10));
+                await Disconnect(ownerUri2);
+            }
+
+            await Task.WhenAll(saveTasks);
+
+            await Disconnect(ownerUri);
+        }
+
+        [Fact]
+        public async Task TestQueryingSubsetDoesNotBlockOtherRequests()
+        {
+            string ownerUri = System.IO.Path.GetTempFileName();
+            string query = "SELECT * FROM sys.objects";
+
+            await Connect(ownerUri, ConnectionTestUtils.AzureTestServerConnection);
+
+            // Execute a query
+            await RunQuery(ownerUri, query);
+
+            // Spawn several tasks for subset requests
+            var subsetTasks = new Task<QueryExecuteSubsetResult>[100];
+            for (int i = 0; i < 100; i++)
+            {
+                subsetTasks[i] = ExecuteSubset(ownerUri, 0, 0, 0, 100);
+            }
+
+            // Interact with the service. None of these requests should time out while waiting for the subset tasks to finish
+            for (int i = 0; i < 10; i++)
+            {
+                string ownerUri2 = System.IO.Path.GetTempFileName();
+
+                await Connect(ownerUri2, ConnectionTestUtils.AzureTestServerConnection);
+                Assert.NotNull(await RequestCompletion(ownerUri2, "SELECT * FROM sys.objects", 0, 10));
+                await Disconnect(ownerUri2);
+            }
+
+            await Task.WhenAll(subsetTasks);
+
+            await Disconnect(ownerUri);
+        }
+
+        [Fact]
+        public async Task TestCancelQueryWhileOtherOperationsAreInProgress()
+        {
+            string ownerUri = System.IO.Path.GetTempFileName();
+            string query = "SELECT * FROM sys.objects a CROSS JOIN sys.objects b";
+            List<Task> tasks = new List<Task>();
+
+            await Connect(ownerUri, ConnectionTestUtils.AzureTestServerConnection);
+
+            // Execute a long-running query
+            var queryTask = RunQuery(ownerUri, query, 60000);
+
+            // Queue up some tasks that interact with the service
+            for (int i = 0; i < 10; i++)
+            {
+                string ownerUri2 = System.IO.Path.GetTempFileName();
+
+                tasks.Add(Task.Run(async () =>
+                {
+                    await Connect(ownerUri2, ConnectionTestUtils.AzureTestServerConnection);
+                    await RequestCompletion(ownerUri2, "SELECT * FROM sys.objects", 0, 10);
+                    await RunQuery(ownerUri2, "SELECT * FROM sys.objects");
+                    await Disconnect(ownerUri2);
+                }));
+            }
+
+            // Cancel the long-running query
+            await CancelQuery(ownerUri);
+
+            await Disconnect(ownerUri);
+        }
+
         [Fact]
         public async Task ExecuteBasicQueryTest()
         {
@@ -78,7 +260,7 @@ namespace Microsoft.SqlTools.ServiceLayer.TestDriver.Tests
             }
         }
 
-        //[Fact]
+        [Fact]
         public async Task TestQueryingAfterCompletionRequests()
         {
             try

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

@@ -211,7 +211,7 @@ namespace Microsoft.SqlTools.ServiceLayer.TestDriver.Tests
         /// <summary>
         /// Run a query using a given connection bound to a URI
         /// </summary>
-        protected async Task<QueryExecuteCompleteParams> RunQuery(string ownerUri, string query)
+        protected async Task<QueryExecuteCompleteParams> RunQuery(string ownerUri, string query, int timeoutMilliseconds = 5000)
         {
             // Write the query text to a backing file
             WriteToFile(ownerUri, query);
@@ -223,7 +223,7 @@ namespace Microsoft.SqlTools.ServiceLayer.TestDriver.Tests
             var result = await Driver.SendRequest(QueryExecuteRequest.Type, queryParams);
             if (result != null && string.IsNullOrEmpty(result.Messages))
             {
-                var eventResult = await Driver.WaitForEvent(QueryExecuteCompleteEvent.Type);
+                var eventResult = await Driver.WaitForEvent(QueryExecuteCompleteEvent.Type, timeoutMilliseconds);
                 return eventResult;
             }
             else
@@ -231,6 +231,64 @@ namespace Microsoft.SqlTools.ServiceLayer.TestDriver.Tests
                 return null;
             }
         }
+        
+        /// <summary>
+        /// Request to cancel an executing query
+        /// </summary>
+        protected async Task<QueryCancelResult> CancelQuery(string ownerUri)
+        {
+            var cancelParams = new QueryCancelParams();
+            cancelParams.OwnerUri = ownerUri;
+
+            var result = await Driver.SendRequest(QueryCancelRequest.Type, cancelParams);
+            return result;
+        }
+
+        /// <summary>
+        /// Request to save query results as CSV
+        /// </summary>
+        protected async Task<SaveResultRequestResult> SaveAsCsv(string ownerUri, string filename, int batchIndex, int resultSetIndex)
+        {
+            var saveParams = new SaveResultsAsCsvRequestParams();
+            saveParams.OwnerUri = ownerUri;
+            saveParams.BatchIndex = batchIndex;
+            saveParams.ResultSetIndex = resultSetIndex;
+            saveParams.FilePath = filename;
+            
+            var result = await Driver.SendRequest(SaveResultsAsCsvRequest.Type, saveParams);
+            return result;
+        }
+
+        /// <summary>
+        /// Request to save query results as JSON
+        /// </summary>
+        protected async Task<SaveResultRequestResult> SaveAsJson(string ownerUri, string filename, int batchIndex, int resultSetIndex)
+        {
+            var saveParams = new SaveResultsAsJsonRequestParams();
+            saveParams.OwnerUri = ownerUri;
+            saveParams.BatchIndex = batchIndex;
+            saveParams.ResultSetIndex = resultSetIndex;
+            saveParams.FilePath = filename;
+            
+            var result = await Driver.SendRequest(SaveResultsAsJsonRequest.Type, saveParams);
+            return result;
+        }
+
+        /// <summary>
+        /// Request a subset of results from a query
+        /// </summary>
+        protected async Task<QueryExecuteSubsetResult> ExecuteSubset(string ownerUri, int batchIndex, int resultSetIndex, int rowStartIndex, int rowCount)
+        {
+            var subsetParams = new QueryExecuteSubsetParams();
+            subsetParams.OwnerUri = ownerUri;
+            subsetParams.BatchIndex = batchIndex;
+            subsetParams.ResultSetIndex = resultSetIndex;
+            subsetParams.RowsStartIndex = rowStartIndex;
+            subsetParams.RowsCount = rowCount;
+
+            var result = await Driver.SendRequest(QueryExecuteSubsetRequest.Type, subsetParams);
+            return result;
+        }
 
         protected void WriteToFile(string ownerUri, string query)
         {