Code Coverage: SaveAsExcel Contracts, EditData Initialization (#463)

* Adding unit tests for save as excel

* Adding unit test for revert cell integration

* Adding unit test for edit session initialization

* Fixing issue where excel file factory wasn't being overridden

test/Microsoft.SqlTools.ServiceLayer.UnitTests/EditData/ServiceIntegrationTests.cs

@@ -4,13 +4,16 @@
 //
 
 using System;
+using System.Data.Common;
 using System.Linq;
 using System.Threading.Tasks;
+using Microsoft.SqlTools.ServiceLayer.Connection;
 using Microsoft.SqlTools.ServiceLayer.EditData;
 using Microsoft.SqlTools.ServiceLayer.EditData.Contracts;
 using Microsoft.SqlTools.ServiceLayer.EditData.UpdateManagement;
 using Microsoft.SqlTools.ServiceLayer.QueryExecution;
 using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts;
+using Microsoft.SqlTools.ServiceLayer.QueryExecution.Contracts.ExecuteRequests;
 using Microsoft.SqlTools.ServiceLayer.Test.Common;
 using Microsoft.SqlTools.ServiceLayer.UnitTests.Utility;
 using Moq;
@@ -158,7 +161,43 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
         }
 
         [Fact]
-        public async Task RevertSucceeds()
+        public async Task RevertCellSucceeds()
+        {
+            // Setup: 
+            // ... Create an edit data service with a session that has a pending cell edit
+            var eds = new EditDataService(null, null, null);
+            var session = await GetDefaultSession();
+            eds.ActiveSessions[Constants.OwnerUri] = session;
+            
+            // ... Make sure that the edit has revert capabilities
+            var mockEdit = new Mock<RowEditBase>();
+            mockEdit.Setup(edit => edit.RevertCell(It.IsAny<int>()))
+                .Returns(new EditRevertCellResult());
+            session.EditCache[0] = mockEdit.Object;
+
+
+            // If: I ask to revert a cell that has a pending edit
+            var efv = new EventFlowValidator<EditRevertCellResult>()
+                .AddResultValidation(Assert.NotNull)
+                .Complete();
+            var param = new EditRevertCellParams
+            {
+                OwnerUri = Constants.OwnerUri,
+                RowId = 0
+            };
+            await eds.HandleRevertCellRequest(param, efv.Object);
+            
+            // Then:
+            // ... It should have succeeded
+            efv.Validate();
+            
+            // ... The edit cache should be empty again
+            EditSession s = eds.ActiveSessions[Constants.OwnerUri];
+            Assert.Empty(s.EditCache);
+        }
+        
+        [Fact]
+        public async Task RevertRowSucceeds()
         {
             // Setup: Create an edit data service with a session that has an pending edit
             var eds = new EditDataService(null, null, null);
@@ -245,6 +284,7 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             efv.Validate();
         }
 
+        #region Initialize Tests
         [Theory]
         [InlineData(null, "table", "table")]            // Null owner URI
         [InlineData(Common.OwnerUri, null, "table")]    // Null object name
@@ -277,11 +317,103 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.EditData
             Assert.Empty(eds.ActiveSessions);
         }
 
+        [Fact]
+        public async Task InitializeSessionExists()
+        {
+            // Setup: Create an edit data service with a session already defined
+            var eds = new EditDataService(null, null, null);
+            var session = await GetDefaultSession();
+            eds.ActiveSessions[Constants.OwnerUri] = session;
+            
+            // If: I request to init a session for an owner URI that already exists
+            var initParams = new EditInitializeParams
+            {
+                ObjectName = "testTable",
+                OwnerUri = Constants.OwnerUri,
+                ObjectType = "Table",
+                Filters = new EditInitializeFiltering()
+            };
+            var efv = new EventFlowValidator<EditInitializeResult>()
+                .AddStandardErrorValidation()
+                .Complete();
+            await eds.HandleInitializeRequest(initParams, efv.Object);
+            
+            // Then:
+            // ... An error event should have been sent
+            efv.Validate();
+            
+            // ... The original session should still be there
+            Assert.Equal(1, eds.ActiveSessions.Count);
+            Assert.Equal(session, eds.ActiveSessions[Constants.OwnerUri]);
+        }
+
+        [Fact]
+        public async Task InitializeSessionSuccess()
+        {
+            // Setup: 
+            // .. Create a mock query
+            var mockQueryResults = QueryExecution.Common.StandardTestDataSet;
+            var cols = mockQueryResults[0].Columns;
+            
+            // ... Create a metadata factory that will return some generic column information
+            var etm = Common.GetStandardMetadata(cols.ToArray());
+            Mock<IEditMetadataFactory> emf = new Mock<IEditMetadataFactory>();
+            emf.Setup(f => f.GetObjectMetadata(It.IsAny<DbConnection>(), It.IsAny<string[]>(), It.IsAny<string>()))
+                .Returns(etm);
+            
+            // ... Create a query execution service that will return a successful query
+            var qes = QueryExecution.Common.GetPrimedExecutionService(mockQueryResults, true, false, null);
+            
+            // ... Create a connection service that doesn't throw when asked for a connection
+            var cs = new Mock<ConnectionService>();
+            cs.Setup(s => s.GetOrOpenConnection(It.IsAny<string>(), It.IsAny<string>(), It.IsAny<bool>()))
+                .Returns(Task.FromResult<DbConnection>(null));
+            
+            // ... Create an edit data service that has mock providers
+            var eds = new EditDataService(qes, cs.Object, emf.Object);
+
+            // If: I request to initialize an edit data session
+            var initParams = new EditInitializeParams
+            {
+                ObjectName = "testTable",
+                OwnerUri = Constants.OwnerUri,
+                ObjectType = "Table",
+                Filters = new EditInitializeFiltering()
+            };
+            var efv = new EventFlowValidator<EditInitializeResult>()
+                .AddResultValidation(Assert.NotNull)
+                .AddEventValidation(BatchStartEvent.Type, Assert.NotNull)
+                .AddEventValidation(ResultSetCompleteEvent.Type, Assert.NotNull)
+                .AddEventValidation(MessageEvent.Type, Assert.NotNull)
+                .AddEventValidation(BatchCompleteEvent.Type, Assert.NotNull)
+                .AddEventValidation(QueryCompleteEvent.Type, Assert.NotNull)
+                .AddEventValidation(EditSessionReadyEvent.Type, esrp =>
+                {
+                    Assert.NotNull(esrp);
+                    Assert.Equal(Constants.OwnerUri, esrp.OwnerUri);
+                    Assert.True(esrp.Success);
+                    Assert.Null(esrp.Message);
+                })
+                .Complete();
+            await eds.HandleInitializeRequest(initParams, efv.Object);
+            await eds.ActiveSessions[Constants.OwnerUri].InitializeTask;
+            
+            // Then:
+            // ... The event should have been received successfully
+            efv.Validate();
+            
+            // ... The session should have been created
+            Assert.Equal(1, eds.ActiveSessions.Count);
+            Assert.True(eds.ActiveSessions.Keys.Contains(Constants.OwnerUri));
+        }
+        
+        #endregion
+        
         [Theory]
-        [InlineData("table", "myschema", new [] { "myschema", "table" })]    // Use schema
-        [InlineData("table", null, new [] { "table" })]    // skip schema
+        [InlineData("table", "myschema", new [] { "myschema", "table" })]                 // Use schema
+        [InlineData("table", null, new [] { "table" })]                                   // skip schema
         [InlineData("schema.table", "myschema", new [] { "myschema", "schema.table"})]    // Use schema
-        [InlineData("schema.table", null, new [] { "schema", "table"})]    // Split object name into schema
+        [InlineData("schema.table", null, new [] { "schema", "table"})]                   // Split object name into schema
         public void ShouldUseSchemaNameIfDefined(string objName, string schemaName, string[] expectedNameParts)
         {
             // Setup: Create an edit data service without a session

test/Microsoft.SqlTools.ServiceLayer.UnitTests/QueryExecution/SaveResults/ServiceIntegrationTests.cs

@@ -137,7 +137,6 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.QueryExecution.SaveResults
 
         [Fact]
         public async Task SaveResultsJsonNonExistentQuery()
-
         {
             // Given: A working query and workspace service
             WorkspaceService<SqlToolsSettings> ws = Common.GetPrimedWorkspaceService(null);
@@ -241,6 +240,115 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.QueryExecution.SaveResults
             efv.Validate();
         }
 
+        #endregion
+        
+        #region Excel Tests 
+        
+        [Fact]
+        public async Task SaveResultsExcelNonExistentQuery()
+        {
+            // Given: A working query and workspace service
+            WorkspaceService<SqlToolsSettings> ws = Common.GetPrimedWorkspaceService(null);
+            QueryExecutionService qes = Common.GetPrimedExecutionService(null, false, false, ws);
+
+            // If: I attempt to save a result set from a query that doesn't exist
+            SaveResultsAsExcelRequestParams saveParams = new SaveResultsAsExcelRequestParams
+            {
+                OwnerUri = Constants.OwnerUri  // Won't exist because nothing has executed
+            };
+            var efv = new EventFlowValidator<SaveResultRequestResult>()
+                .AddStandardErrorValidation()
+                .Complete();
+            await qes.HandleSaveResultsAsExcelRequest(saveParams, efv.Object);
+
+            // Then:
+            // ... An error event should have been fired
+            // ... No success event should have been fired
+            efv.Validate();
+        }
+
+        [Fact]
+        public async Task SaveResultAsExcelFailure()
+        {
+            // Given: 
+            // ... A working query and workspace service
+            WorkspaceService<SqlToolsSettings> ws = Common.GetPrimedWorkspaceService(Constants.StandardQuery);
+            Dictionary<string, byte[]> storage;
+            QueryExecutionService qes = Common.GetPrimedExecutionService(Common.StandardTestDataSet, true, false, ws, out storage);
+
+            // ... The query execution service has executed a query with results
+            var executeParams = new ExecuteDocumentSelectionParams { QuerySelection = null, OwnerUri = Constants.OwnerUri };
+            var executeRequest = RequestContextMocks.Create<ExecuteRequestResult>(null);
+            await qes.HandleExecuteRequest(executeParams, executeRequest.Object);
+            await qes.ActiveQueries[Constants.OwnerUri].ExecutionTask;
+
+            // If: I attempt to save a result set and get it to throw because of invalid column selection
+            SaveResultsAsExcelRequestParams saveParams = new SaveResultsAsExcelRequestParams
+            {
+                BatchIndex = 0,
+                FilePath = "qqq",
+                OwnerUri = Constants.OwnerUri,
+                ResultSetIndex = 0,
+                ColumnStartIndex = -1,
+                ColumnEndIndex = 100,
+                RowStartIndex = 0,
+                RowEndIndex = 5
+            };
+            qes.JsonFileFactory = GetExcelStreamFactory(storage, saveParams);
+            var efv = new EventFlowValidator<SaveResultRequestResult>()
+                .AddStandardErrorValidation()
+                .Complete();
+            await qes.HandleSaveResultsAsExcelRequest(saveParams, efv.Object);
+            await qes.ActiveQueries[saveParams.OwnerUri]
+                .Batches[saveParams.BatchIndex]
+                .ResultSets[saveParams.ResultSetIndex]
+                .SaveTasks[saveParams.FilePath];
+
+            // Then:
+            // ... An error event should have been fired
+            // ... No success event should have been fired
+            efv.Validate();
+        }
+
+        [Fact]
+        public async Task SaveResultsAsExcelSuccess()
+        {
+            // Given: 
+            // ... A working query and workspace service
+            WorkspaceService<SqlToolsSettings> ws = Common.GetPrimedWorkspaceService(Constants.StandardQuery);
+            Dictionary<string, byte[]> storage;
+            QueryExecutionService qes = Common.GetPrimedExecutionService(Common.StandardTestDataSet, true, false, ws, out storage);
+
+            // ... The query execution service has executed a query with results
+            var executeParams = new ExecuteDocumentSelectionParams { QuerySelection = null, OwnerUri = Constants.OwnerUri };
+            var executeRequest = RequestContextMocks.Create<ExecuteRequestResult>(null);
+            await qes.HandleExecuteRequest(executeParams, executeRequest.Object);
+            await qes.ActiveQueries[Constants.OwnerUri].ExecutionTask;
+
+            // If: I attempt to save a result set from a query
+            SaveResultsAsExcelRequestParams saveParams = new SaveResultsAsExcelRequestParams
+            {
+                OwnerUri = Constants.OwnerUri,
+                FilePath = "qqq",
+                BatchIndex = 0,
+                ResultSetIndex = 0
+            };
+            qes.ExcelFileFactory = GetExcelStreamFactory(storage, saveParams);
+            var efv = new EventFlowValidator<SaveResultRequestResult>()
+                .AddStandardResultValidator()
+                .Complete();
+            await qes.HandleSaveResultsAsExcelRequest(saveParams, efv.Object);
+            await qes.ActiveQueries[saveParams.OwnerUri]
+                .Batches[saveParams.BatchIndex]
+                .ResultSets[saveParams.ResultSetIndex]
+                .SaveTasks[saveParams.FilePath];
+
+            // Then:
+            // ... I should have a successful result
+            // ... There should not have been an error
+            efv.Validate();
+        }
+        
         #endregion
 
         #region Private Helpers
@@ -275,6 +383,22 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.QueryExecution.SaveResults
             return mock.Object;
         }
 
+        private static IFileStreamFactory GetExcelStreamFactory(IDictionary<string, byte[]> storage,
+            SaveResultsAsExcelRequestParams saveParams)
+        {
+            Mock<IFileStreamFactory> mock = new Mock<IFileStreamFactory>();
+            mock.Setup(fsf => fsf.GetReader(It.IsAny<string>()))
+                .Returns<string>(output => new ServiceBufferFileStreamReader(new MemoryStream(storage[output]), new QueryExecutionSettings()));
+            mock.Setup(fsf => fsf.GetWriter(It.IsAny<string>()))
+                .Returns<string>(output =>
+                {
+                    storage.Add(output, new byte[8192]);
+                    return new SaveAsExcelFileStreamWriter(new MemoryStream(storage[output]), saveParams);
+                });
+            
+            return mock.Object;
+        }
+
         #endregion
     }