added handler for oe refresh and close session requests (#332)

* added handler for oe refresh and close session requests

src/Microsoft.SqlTools.ServiceLayer/Connection/ConnectionService.cs

@@ -476,7 +476,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Connection
         /// <summary>
         /// Close a connection with the specified connection details.
         /// </summary>
-        public bool Disconnect(DisconnectParams disconnectParams)
+        public virtual bool Disconnect(DisconnectParams disconnectParams)
         {
             // Validate parameters
             if (disconnectParams == null || string.IsNullOrEmpty(disconnectParams.OwnerUri))

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Contracts/CloseSessionRequest.cs

@@ -0,0 +1,53 @@
+//
+// 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.Hosting.Protocol.Contracts;
+using Microsoft.SqlTools.ServiceLayer.Connection.Contracts;
+
+namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Contracts
+{
+    /// <summary>
+    /// Information returned from a <see cref="CloseSessionRequest"/>.
+    /// Contains success information, a <see cref="SessionId"/> to be used when
+    /// requesting closing an existing session.
+    /// </summary>
+    public class CloseSessionResponse
+    {
+        /// <summary>
+        /// Boolean indicating if the session was closed successfully 
+        /// </summary>
+        public bool Success { get; set; }
+
+        /// <summary>
+        /// Unique ID to use when sending any requests for objects in the
+        /// tree under the node
+        /// </summary>
+        public string SessionId { get; set; }
+    }
+
+    /// <summary>
+    /// Parameters to the <see cref="CloseSessionRequest"/>.
+    /// </summary>
+    public class CloseSessionParams
+    {
+        /// <summary>
+        /// The Id returned from a <see cref="CreateSessionRequest"/>. This
+        /// is used to disambiguate between different trees. 
+        /// </summary>
+        public string SessionId { get; set; }
+    }
+
+    /// <summary>
+    /// Establishes an Object Explorer tree session for a specific connection.
+    /// This will create a connection to a specific server or database, register
+    /// it for use in the 
+    /// </summary>
+    public class CloseSessionRequest
+    {
+        public static readonly
+            RequestType<CloseSessionParams, CloseSessionResponse> Type =
+            RequestType<CloseSessionParams, CloseSessionResponse>.Create("objectexplorer/closesession");
+    }
+}

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Contracts/RefreshRequest.cs

@@ -7,6 +7,23 @@ using Microsoft.SqlTools.Hosting.Protocol.Contracts;
 
 namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Contracts
 {
+    /// <summary>
+    /// Information returned from a <see cref="RefreshRequest"/>.
+    /// </summary>
+    public class RefreshResponse
+    {
+        /// <summary>
+        /// Unique ID to use when sending any requests for objects in the
+        /// tree under the node
+        /// </summary>
+        public string SessionId { get; set; }
+
+        /// <summary>
+        /// Information describing the expanded nodes in the tree
+        /// </summary>
+        public NodeInfo[] Nodes { get; set; }
+    }
+
     /// <summary>
     /// Parameters to the <see cref="ExpandRequest"/>.
     /// </summary>
@@ -19,9 +36,9 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Contracts
         public string SessionId { get; set; }
 
         /// <summary>
-        /// Path identifying the node to expand. See <see cref="NodeInfo.NodePath"/> for details
+        /// Path identifying the node to refresh. See <see cref="NodeInfo.NodePath"/> for details
         /// </summary>
-        public string[] NodePath { get; set; }
+        public string NodePath { get; set; }
     }
 
     /// <summary>
@@ -33,7 +50,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Contracts
         /// Returns children of a given node as a <see cref="NodeInfo"/> array.
         /// </summary>
         public static readonly
-            RequestType<ExpandParams, NodeInfo[]> Type =
+            RequestType<RefreshParams, RefreshResponse> Type =
-            RequestType<ExpandParams, NodeInfo[]>.Create("objectexplorer/refresh");
+            RequestType<RefreshParams, RefreshResponse>.Create("objectexplorer/refresh");
     }
 }

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Nodes/NodeFilter.cs

@@ -74,7 +74,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes
                     {
                         proeprtyValue = (int)Convert.ChangeType(value, Type);
                     }
-                    string orPrefix = i == 0 ? "" : "or";
+                    string orPrefix = i == 0 ? string.Empty : "or";
                     filter = $"{filter} {orPrefix} @{Property} = {proeprtyValue}";
                 }
             }
@@ -91,8 +91,8 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes
             {
                 var value = list[i];
                 
-                string orPrefix = i == 0 ? "" : "and";
+                string andPrefix = i == 0 ? string.Empty : "and";
-                filter = $"{filter} {orPrefix} {value.ToPropertyFilterString()}";
+                filter = $"{filter} {andPrefix} {value.ToPropertyFilterString()}";
             }
             filter = $"[{filter}]";
 

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/Nodes/TreeNode.cs

@@ -195,6 +195,17 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes
             return children;
         }
 
+        /// <summary>
+        /// Refresh this node and returns its children
+        /// </summary>
+        /// <returns>Children as an IList. This is the raw children collection, not a copy</returns>
+        public IList<TreeNode> Refresh()
+        {
+            // TODO consider why solution explorer has separate Children and Items options
+            PopulateChildren();
+            return children;
+        }
+
         /// <summary>
         /// Gets a readonly view of the currently defined children for this node. 
         /// This does not expand the node at all

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/ObjectExplorerService.cs

@@ -4,7 +4,9 @@
 //
 
 using System;
+using System.Collections.Concurrent;
 using System.Collections.Generic;
+using System.Collections.ObjectModel;
 using System.Composition;
 using System.Globalization;
 using System.Linq;
@@ -33,7 +35,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
         // Instance of the connection service, used to get the connection info for a given owner URI
         private ConnectionService connectionService;
         private IProtocolEndpoint serviceHost;
-        private Dictionary<string, ObjectExplorerSession> sessionMap;
+        private ConcurrentDictionary<string, ObjectExplorerSession> sessionMap;
         private readonly Lazy<Dictionary<string, HashSet<ChildFactory>>> applicableNodeChildFactories;
         private IMultiServiceProvider serviceProvider;
 
@@ -42,7 +44,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
         /// </summary>
         public ObjectExplorerService()
         {
-            sessionMap = new Dictionary<string, ObjectExplorerSession>();
+            sessionMap = new ConcurrentDictionary<string, ObjectExplorerSession>();
             applicableNodeChildFactories = new Lazy<Dictionary<string, HashSet<ChildFactory>>>(() => PopulateFactories());
         }
 
@@ -63,6 +65,17 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
             }
         }
 
+        /// <summary>
+        /// Returns the session ids
+        /// </summary>
+        internal IReadOnlyCollection<string> SessionIds
+        {
+            get
+            {
+                return new ReadOnlyCollection<string>(sessionMap.Keys.ToList());
+            }
+        } 
+    
         /// <summary>
         /// As an <see cref="IComposableService"/>, this will be set whenever the service is initialized
         /// via an <see cref="IMultiServiceProvider"/>
@@ -86,6 +99,8 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
             // Register handlers for requests
             serviceHost.SetRequestHandler(CreateSessionRequest.Type, HandleCreateSessionRequest);
             serviceHost.SetRequestHandler(ExpandRequest.Type, HandleExpandRequest);
+            serviceHost.SetRequestHandler(RefreshRequest.Type, HandleRefreshRequest);
+            serviceHost.SetRequestHandler(CloseSessionRequest.Type, HandleCloseSessionRequest);
         }
 
         public void CloseSession(string uri)
@@ -94,7 +109,11 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
             if (sessionMap.TryGetValue(uri, out session))
             {
                 // Establish a connection to the specified server/database
-                sessionMap.Remove(session.Uri);
+                sessionMap.TryRemove(session.Uri, out session);
+                connectionService.Disconnect(new DisconnectParams()
+                {
+                    OwnerUri = uri
+                });
             }
         }
         
@@ -136,16 +155,23 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
             await HandleRequestAsync(doCreateSession, context, "HandleCreateSessionRequest");
         }
 
-        internal async Task<NodeInfo[]> ExpandNode(ObjectExplorerSession session, string nodePath)
+        internal async Task<NodeInfo[]> ExpandNode(ObjectExplorerSession session, string nodePath, bool forceRefresh = false)
         {
             return await Task.Factory.StartNew(() =>
             {
                 NodeInfo[] nodes = null;
                 TreeNode node = session.Root.FindNodeByPath(nodePath);
                 if(node != null)
+                {
+                    if (forceRefresh)
+                    {
+                        nodes = node.Refresh().Select(x => x.ToNodeInfo()).ToArray();
+                    }
+                    else
                     {
                         nodes = node.Expand().Select(x => x.ToNodeInfo()).ToArray();
                     }
+                }
                 return nodes; 
             });
         }
@@ -167,7 +193,7 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
             }
 
             session = ObjectExplorerSession.CreateSession(connectionResult, serviceProvider);
-            sessionMap[uri] = session;
+            sessionMap.AddOrUpdate(uri, session, (key, oldSession) => session);
             return session;
         }
         
@@ -212,23 +238,18 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
                 string uri = expandParams.SessionId;
                 ObjectExplorerSession session = null;
                 NodeInfo[] nodes = null;
-                if (sessionMap.ContainsKey(uri))
+                ExpandResponse response;
+                if (!sessionMap.TryGetValue(uri, out session))
                 {
-                    session = sessionMap[uri];
+                    Logger.Write(LogLevel.Verbose, $"Cannot expand object explorer node. Couldn't find session for uri. {uri} ");
                 }
-                else
-                {
-                    //TODO: error
-                }
-
 
                 if (session != null)
                 {
-                    // Establish a connection to the specified server/database
+                    // expand the nodes for given node path
                     nodes = await ExpandNode(session, expandParams.NodePath);
                 }
 
-                ExpandResponse response;
                 response = new ExpandResponse() { Nodes = nodes, SessionId = uri };
                 return response;
             };
@@ -236,6 +257,68 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer
             await HandleRequestAsync(expandNode, context, "HandleExpandRequest");
         }
 
+        internal async Task HandleRefreshRequest(RefreshParams refreshParams, RequestContext<RefreshResponse> context)
+        {
+            Logger.Write(LogLevel.Verbose, "HandleRefreshRequest");
+            Func<Task<RefreshResponse>> refreshNode = async () =>
+            {
+                Validate.IsNotNull(nameof(refreshParams), refreshParams);
+                Validate.IsNotNull(nameof(context), context);
+
+                string uri = refreshParams.SessionId;
+                ObjectExplorerSession session = null;
+                NodeInfo[] nodes = null;
+                RefreshResponse response;
+                if (!sessionMap.TryGetValue(uri, out session))
+                {
+                    Logger.Write(LogLevel.Verbose, $"Cannot refresh object explorer node. Couldn't find session for uri. {uri} ");
+                }
+
+                if (session != null)
+                {
+                    // refresh the nodes for given node path
+                    nodes = await ExpandNode(session, refreshParams.NodePath, true);
+                }
+
+                response = new RefreshResponse() { Nodes = nodes, SessionId = uri };
+                return response;
+            };
+
+            await HandleRequestAsync(refreshNode, context, "HandleRefreshRequest");
+        }
+
+        internal async Task HandleCloseSessionRequest(CloseSessionParams closeSessionParams, RequestContext<CloseSessionResponse> context)
+        {
+            Validate.IsNotNull(nameof(closeSessionParams), closeSessionParams);
+            Validate.IsNotNull(nameof(context), context);
+            Logger.Write(LogLevel.Verbose, "HandleCloseSessionRequest");
+            Func<Task<CloseSessionResponse>> closeSession = () =>
+            {
+                return Task.Factory.StartNew(() =>
+                {
+                    string uri = closeSessionParams.SessionId;
+                    ObjectExplorerSession session = null;
+                    bool success = false;
+                    if (!sessionMap.TryGetValue(uri, out session))
+                    {
+                        Logger.Write(LogLevel.Verbose, $"Cannot close object explorer session. Couldn't find session for uri. {uri} ");
+                    }
+
+                    if (session != null)
+                    {
+                        // refresh the nodes for given node path
+                        CloseSession(uri);
+                        success = true;
+                    }
+
+                    var response = new CloseSessionResponse() {Success = success, SessionId = uri};
+                    return response;
+                });
+            };
+
+            await HandleRequestAsync(closeSession, context, "HandleCloseSessionRequest");
+        }
+
         private async Task HandleRequestAsync<T>(Func<Task<T>> handler, RequestContext<T> requestContext, string requestType)
         {
             Logger.Write(LogLevel.Verbose, requestType);

src/Microsoft.SqlTools.ServiceLayer/ObjectExplorer/SmoModel/SmoQueryModel.tt

@@ -64,10 +64,18 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
 			WriteLine(string.Format("if ({0} != null)", parentVar));
 			WriteLine("{");
 			PushIndent(indent);
+			WriteLine("bool hasFilter = !string.IsNullOrEmpty(filter);");
 			
 			string navigationPath = GetNavigationPath(nodeElement, xmlFile, nodeName, parentType);
 			WriteLine(string.Format("var retValue = {0}.{1};", parentVar, navigationPath));
-			WriteLine("bool hasFilter = !string.IsNullOrEmpty(filter);");
+			WriteLine("if (retValue != null)");
+			WriteLine("{");
+			PushIndent(indent);
+			
+			
+			if (IsCollection(nodeElement))
+			{
+				
 				WriteLine("HashSet<string> urns = null;");
 				
 				WriteLine("if (hasFilter)");
@@ -81,11 +89,6 @@ namespace Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel
 				WriteLine("urns = GetUrns(result);");
 				PopIndent();
 				WriteLine("}");
-			WriteLine("if (retValue != null)");
-			WriteLine("{");
-			PushIndent(indent);
-			if (IsCollection(nodeElement))
-			{
 				WriteLine("if (hasFilter && urns != null)");
 				WriteLine("{");
 				PushIndent(indent);

test/CodeCoverage/gulpfile.js

@@ -92,7 +92,7 @@ gulp.task('ext:nuget-restore', function() {
 
 
 gulp.task('ext:code-coverage', function(done) {
-    cproc.execFile('cmd.exe', [ '/c', 'codecoverage.bat' ], function(err, stdout) {
+    cproc.execFile('cmd.exe', [ '/c', 'codecoverage.bat' ], {maxBuffer: 1024 * 500}, function(err, stdout) {
         if (err) {
             throw new gutil.PluginError('ext:code-coverage', err);
         }

test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/BatchParser/BatchParserTests.cs

@@ -252,7 +252,7 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.BatchParser
 
             string outputString = output.ToString().Replace("\r\n", "\n");
 
-            Console.WriteLine(baselineFilename);
+            //Console.WriteLine(baselineFilename);
 
             if (string.Compare(baseline, outputString, StringComparison.Ordinal) != 0)
             {

test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/LanguageServer/PeekDefinitionTests.cs

@@ -19,6 +19,8 @@ using Xunit;
 using ConnectionType = Microsoft.SqlTools.ServiceLayer.Connection.ConnectionType;
 using Location = Microsoft.SqlTools.ServiceLayer.Workspace.Contracts.Location;
 using System.Collections.Generic;
+using System.Threading.Tasks;
+using Microsoft.SqlTools.ServiceLayer.Connection.Contracts;
 
 namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.LanguageServices
 {
@@ -348,19 +350,18 @@ GO";
         /// Test get definition for a scalar valued function object with active connection and explicit schema name. Expect non-null locations
         /// </summary>
         [Fact]
-        public void GetScalarValuedFunctionDefinitionWithSchemaNameSuccessTest()
+        public async Task GetScalarValuedFunctionDefinitionWithSchemaNameSuccessTest()
         {
-            ExecuteAndValidatePeekTest(AddTwoFunctionQuery, AddTwoFunctionName, ScalarValuedFunctionTypeName);
+            await ExecuteAndValidatePeekTest(AddTwoFunctionQuery, AddTwoFunctionName, ScalarValuedFunctionTypeName);
         }
 
-        private void ExecuteAndValidatePeekTest(string query, string objectName, string objectType, string schemaName = "dbo")
+        private async Task ExecuteAndValidatePeekTest(string query, string objectName, string objectType, string schemaName = "dbo")
         {
             if (!string.IsNullOrEmpty(query))
             {
-                using (SqlTestDb testDb = SqlTestDb.CreateNew(TestServerType.OnPrem, query))
+                SqlTestDb testDb = await SqlTestDb.CreateNewAsync(TestServerType.OnPrem, false, null, query);
-                {
                 ValidatePeekTest(testDb.DatabaseName, objectName, objectType, schemaName, true);
-                }
+                await testDb.CleanupAsync();
             }
             else
             {
@@ -419,58 +420,64 @@ GO";
             {
                 Assert.Null(locations);
             }
+
+            var connectionService = LiveConnectionHelper.GetLiveTestConnectionService();
+            connectionService.Disconnect(new DisconnectParams
+            {
+                    OwnerUri = connInfo.OwnerUri
+            });
         }
 
         /// <summary>
         /// Test get definition for a table valued function object with active connection and explicit schema name. Expect non-null locations
         /// </summary>
         [Fact]
-        public void GetTableValuedFunctionDefinitionWithSchemaNameSuccessTest()
+        public async Task GetTableValuedFunctionDefinitionWithSchemaNameSuccessTest()
         {
-            ExecuteAndValidatePeekTest(ReturnTableTableFunctionQuery, ReturnTableFunctionName, TableValuedFunctionTypeName);
+            await ExecuteAndValidatePeekTest(ReturnTableTableFunctionQuery, ReturnTableFunctionName, TableValuedFunctionTypeName);
         }
 
         /// <summary>
         /// Test get definition for a scalar valued function object that doesn't exist with active connection. Expect null locations
         /// </summary>
         [Fact]
-        public void GetScalarValuedFunctionDefinitionWithNonExistentFailureTest()
+        public async Task GetScalarValuedFunctionDefinitionWithNonExistentFailureTest()
         {
             string objectName = "doesNotExist";
             string schemaName = "dbo";
             string objectType = ScalarValuedFunctionTypeName;
 
-            ExecuteAndValidatePeekTest(null, objectName, objectType, schemaName);
+            await ExecuteAndValidatePeekTest(null, objectName, objectType, schemaName);
         }
 
         /// <summary>
         /// Test get definition for a table valued function object that doesn't exist with active connection. Expect null locations
         /// </summary>
         [Fact]
-        public void GetTableValuedFunctionDefinitionWithNonExistentObjectFailureTest()
+        public async Task GetTableValuedFunctionDefinitionWithNonExistentObjectFailureTest()
         {
             string objectName = "doesNotExist";
             string schemaName = "dbo";
             string objectType = TableValuedFunctionTypeName;
-            ExecuteAndValidatePeekTest(null, objectName, objectType, schemaName);
+            await ExecuteAndValidatePeekTest(null, objectName, objectType, schemaName);
         }
 
         /// <summary>
         /// Test get definition for a scalar valued function object with active connection. Expect non-null locations
         /// </summary>
         [Fact]
-        public void GetScalarValuedFunctionDefinitionWithoutSchemaNameSuccessTest()
+        public async Task GetScalarValuedFunctionDefinitionWithoutSchemaNameSuccessTest()
         {
-            ExecuteAndValidatePeekTest(AddTwoFunctionQuery, AddTwoFunctionName, ScalarValuedFunctionTypeName, null);
+            await ExecuteAndValidatePeekTest(AddTwoFunctionQuery, AddTwoFunctionName, ScalarValuedFunctionTypeName, null);
         }
 
         /// <summary>
         /// Test get definition for a table valued function object with active connection. Expect non-null locations
         /// </summary>
         [Fact]
-        public void GetTableValuedFunctionDefinitionWithoutSchemaNameSuccessTest()
+        public async Task GetTableValuedFunctionDefinitionWithoutSchemaNameSuccessTest()
         {
-            ExecuteAndValidatePeekTest(ReturnTableTableFunctionQuery, ReturnTableFunctionName, TableValuedFunctionTypeName, null);
+            await ExecuteAndValidatePeekTest(ReturnTableTableFunctionQuery, ReturnTableFunctionName, TableValuedFunctionTypeName, null);
         }
 
 
@@ -478,60 +485,60 @@ GO";
         /// Test get definition for a user defined data type object with active connection and explicit schema name. Expect non-null locations
         /// </summary>
         [Fact]
-        public void GetUserDefinedDataTypeDefinitionWithSchemaNameSuccessTest()
+        public async Task GetUserDefinedDataTypeDefinitionWithSchemaNameSuccessTest()
         {
-            ExecuteAndValidatePeekTest(SsnTypeQuery, SsnTypeName, UserDefinedDataTypeTypeName);
+            await ExecuteAndValidatePeekTest(SsnTypeQuery, SsnTypeName, UserDefinedDataTypeTypeName);
         }
 
         /// <summary>
         /// Test get definition for a user defined data type object with active connection. Expect non-null locations
         /// </summary>
         [Fact]
-        public void GetUserDefinedDataTypeDefinitionWithoutSchemaNameSuccessTest()
+        public async Task GetUserDefinedDataTypeDefinitionWithoutSchemaNameSuccessTest()
         {
-            ExecuteAndValidatePeekTest(SsnTypeQuery, SsnTypeName, UserDefinedDataTypeTypeName, null);
+            await ExecuteAndValidatePeekTest(SsnTypeQuery, SsnTypeName, UserDefinedDataTypeTypeName, null);
         }
 
         /// <summary>
         /// Test get definition for a user defined data type object that doesn't exist with active connection. Expect null locations
         /// </summary>
         [Fact]
-        public void GetUserDefinedDataTypeDefinitionWithNonExistentFailureTest()
+        public async Task GetUserDefinedDataTypeDefinitionWithNonExistentFailureTest()
         {
             string objectName = "doesNotExist";
             string schemaName = "dbo";
             string objectType = UserDefinedDataTypeTypeName;
-            ExecuteAndValidatePeekTest(null, objectName, objectType, schemaName);
+            await ExecuteAndValidatePeekTest(null, objectName, objectType, schemaName);
         }
 
         /// <summary>
         /// Test get definition for a user defined table type object with active connection and explicit schema name. Expect non-null locations
         /// </summary>
         [Fact]
-        public void GetUserDefinedTableTypeDefinitionWithSchemaNameSuccessTest()
+        public async Task GetUserDefinedTableTypeDefinitionWithSchemaNameSuccessTest()
         {
-            ExecuteAndValidatePeekTest(LocationTableTypeQuery, LocationTableTypeName, UserDefinedTableTypeTypeName);
+            await ExecuteAndValidatePeekTest(LocationTableTypeQuery, LocationTableTypeName, UserDefinedTableTypeTypeName);
         }
 
         /// <summary>
         /// Test get definition for a user defined table type object with active connection. Expect non-null locations
         /// </summary>
         [Fact]
-        public void GetUserDefinedTableTypeDefinitionWithoutSchemaNameSuccessTest()
+        public async Task GetUserDefinedTableTypeDefinitionWithoutSchemaNameSuccessTest()
         {
-            ExecuteAndValidatePeekTest(LocationTableTypeQuery, LocationTableTypeName, UserDefinedTableTypeTypeName, null);
+            await ExecuteAndValidatePeekTest(LocationTableTypeQuery, LocationTableTypeName, UserDefinedTableTypeTypeName, null);
         }
 
         /// <summary>
         /// Test get definition for a user defined table type object that doesn't exist with active connection. Expect null locations
         /// </summary>
         [Fact]
-        public void GetUserDefinedTableTypeDefinitionWithNonExistentFailureTest()
+        public async Task GetUserDefinedTableTypeDefinitionWithNonExistentFailureTest()
         {
             string objectName = "doesNotExist";
             string schemaName = "dbo";
             string objectType = UserDefinedTableTypeTypeName;
-            ExecuteAndValidatePeekTest(null, objectName, objectType, schemaName);
+            await ExecuteAndValidatePeekTest(null, objectName, objectType, schemaName);
 
         }
 
@@ -539,9 +546,9 @@ GO";
         /// Test get definition for a synonym object with active connection and explicit schema name. Expect non-null locations
         /// </summary>
         [Fact]
-        public void GetSynonymDefinitionWithSchemaNameSuccessTest()
+        public async Task GetSynonymDefinitionWithSchemaNameSuccessTest()
         {
-            ExecuteAndValidatePeekTest(TestTableSynonymQuery, TestTableSynonymName, SynonymTypeName);
+            await ExecuteAndValidatePeekTest(TestTableSynonymQuery, TestTableSynonymName, SynonymTypeName);
         }
 
 
@@ -549,21 +556,21 @@ GO";
         /// Test get definition for a Synonym object with active connection. Expect non-null locations
         /// </summary>
         [Fact]
-        public void GetSynonymDefinitionWithoutSchemaNameSuccessTest()
+        public async Task GetSynonymDefinitionWithoutSchemaNameSuccessTest()
         {
-            ExecuteAndValidatePeekTest(TestTableSynonymQuery, TestTableSynonymName, SynonymTypeName, null);
+            await ExecuteAndValidatePeekTest(TestTableSynonymQuery, TestTableSynonymName, SynonymTypeName, null);
         }
 
         /// <summary>
         /// Test get definition for a Synonym object that doesn't exist with active connection. Expect null locations
         /// </summary>
         [Fact]
-        public void GetSynonymDefinitionWithNonExistentFailureTest()
+        public async Task GetSynonymDefinitionWithNonExistentFailureTest()
         {
             string objectName = "doesNotExist";
             string schemaName = "dbo";
             string objectType = "Synonym";
-            ExecuteAndValidatePeekTest(null, objectName, objectType, schemaName);
+            await ExecuteAndValidatePeekTest(null, objectName, objectType, schemaName);
         }
 
         /// <summary>
@@ -701,7 +708,7 @@ GO";
         /// Get Definition for a object by putting the cursor on 3 different
         /// objects
         /// </summary>
-        [Fact]
+       // [Fact]
         public async void GetDefinitionFromChildrenAndParents()
         {
             string queryString = "select * from master.sys.objects";
@@ -758,7 +765,7 @@ GO";
             connInfo.RemoveAllConnections();
         }
 
-        [Fact]
+       // [Fact]
         public async void GetDefinitionFromProcedures()
         {
 

test/Microsoft.SqlTools.ServiceLayer.IntegrationTests/ObjectExplorer/ObjectExplorerServiceTests.cs

@@ -9,9 +9,7 @@ using System.Globalization;
 using System.IO;
 using System.Linq;
 using System.Text;
-using System.Threading;
 using System.Threading.Tasks;
-using Microsoft.SqlTools.ServiceLayer.Connection;
 using Microsoft.SqlTools.ServiceLayer.Connection.Contracts;
 using Microsoft.SqlTools.ServiceLayer.ObjectExplorer;
 using Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Contracts;
@@ -31,52 +29,42 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.ObjectExplorer
         public async void CreateSessionAndExpandOnTheServerShouldReturnServerAsTheRoot()
         {
             var query = "";
-            string uri = "CreateSessionAndExpandServer";
             string databaseName = null;
-            using (SqlTestDb testDb = SqlTestDb.CreateNew(TestServerType.OnPrem, false, databaseName, query, uri))
+            await RunTest(databaseName, query, "EmptyDatabase", async (testDbName, session) =>
             {
-                var session = await CreateSession(null, uri);
+                await ExpandServerNodeAndVerifyDatabaseHierachy(testDbName, session);
-                await ExpandServerNodeAndVerifyDatabaseHierachy(testDb.DatabaseName, session);
+            });
-                DisconnectConnection(uri);
-            }
         }
 
         [Fact]
         public async void CreateSessionWithTempdbAndExpandOnTheServerShouldReturnServerAsTheRoot()
         {
             var query = "";
-            string uri = "CreateSessionAndExpandServer";
+            string databaseName = "tempdb";
-            string databaseName = null;
+            await RunTest(databaseName, query, "TepmDb", async (testDbName, session) =>
-            using (SqlTestDb testDb = SqlTestDb.CreateNew(TestServerType.OnPrem, false, databaseName, query, uri))
             {
-                var session = await CreateSession("tempdb", uri);
+                await ExpandServerNodeAndVerifyDatabaseHierachy(testDbName, session);
-                await ExpandServerNodeAndVerifyDatabaseHierachy(testDb.DatabaseName, session);
+            });
-                DisconnectConnection(uri);
-            }
         }
 
         [Fact]
         public async void CreateSessionAndExpandOnTheDatabaseShouldReturnDatabaseAsTheRoot()
         {
             var query = "";
-            string uri = "CreateSessionAndExpandDatabase";
+            string databaseName = "#testDb#";
-            string databaseName = null;
+            await RunTest(databaseName, query, "TestDb", async (testDbName, session) =>
-            using (SqlTestDb testDb = SqlTestDb.CreateNew(TestServerType.OnPrem, false, databaseName, query, uri))
             {
-                var session = await CreateSession(testDb.DatabaseName, uri);
+                await ExpandAndVerifyDatabaseNode(testDbName, session);
-                ExpandAndVerifyDatabaseNode(testDb.DatabaseName, session);
+            });
-                DisconnectConnection(uri);
-            }
         }
 
         [Fact]
         public async void VerifyAllSqlObjects()
         {
             var queryFileName = "AllSqlObjects.sql";
-            string uri = "AllSqlObjects";
             string baselineFileName = "AllSqlObjects.txt";
-            string databaseName = null;
+            string databaseName = "#testDb#";
-            await TestServiceProvider.CalculateRunTime(() => VerifyObjectExplorerTest(databaseName, queryFileName, uri, baselineFileName), true);
+            await TestServiceProvider.CalculateRunTime(() => VerifyObjectExplorerTest(databaseName, "AllSqlObjects", queryFileName, baselineFileName), true);
         }
 
         //[Fact]
@@ -84,17 +72,50 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.ObjectExplorer
         public async void VerifySystemObjects()
         {
             string queryFileName = null;
-            string uri = "SystemObjects";
             string baselineFileName = null;
-            string databaseName = null;
+            string databaseName = "#testDb#";
-            await TestServiceProvider.CalculateRunTime(() => VerifyObjectExplorerTest(databaseName, queryFileName, uri, baselineFileName, true), true);
+            await TestServiceProvider.CalculateRunTime(() => VerifyObjectExplorerTest(databaseName, queryFileName, "SystemOBjects", baselineFileName, true), true);
         }
 
-        private async Task<ObjectExplorerSession> CreateSession(string databaseName, string uri)
+        private async Task RunTest(string databaseName, string query, string testDbPrefix, Func<string, ObjectExplorerSession, Task> test)
+        {
+            SqlTestDb testDb = null;
+            string uri = string.Empty;
+            try
+            {
+                testDb = await SqlTestDb.CreateNewAsync(TestServerType.OnPrem, false, null, query, testDbPrefix);
+                if (databaseName == "#testDb#")
+                {
+                    databaseName = testDb.DatabaseName;
+                }
+                var session = await CreateSession(databaseName);
+                uri = session.Uri;
+                await test(testDb.DatabaseName, session);
+            }
+            catch (Exception ex)
+            {
+                Console.WriteLine($"Failed to run OE test. uri:{uri} error:{ex.Message}");
+                Assert.False(true, ex.Message);
+            }
+            finally
+            {
+                if (!string.IsNullOrEmpty(uri))
+                {
+                    CloseSession(uri);
+                }
+                if (testDb != null)
+                {
+                    await testDb.CleanupAsync();
+                }
+            }
+        }
+
+        private async Task<ObjectExplorerSession> CreateSession(string databaseName)
         {
             ConnectParams connectParams = TestServiceProvider.Instance.ConnectionProfileService.GetConnectionParameters(TestServerType.OnPrem, databaseName);
             //connectParams.Connection.Pooling = false;
             ConnectionDetails details = connectParams.Connection;
+            string uri = ObjectExplorerService.GenerateUri(details);
 
             var session =  await _service.DoCreateSession(details, uri);
             Console.WriteLine(string.Format(CultureInfo.InvariantCulture, "OE session created for database: {0}", databaseName));
@@ -139,7 +160,7 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.ObjectExplorer
             return databaseNode;
         }
 
-        private void ExpandAndVerifyDatabaseNode(string databaseName, ObjectExplorerSession session)
+        private async Task ExpandAndVerifyDatabaseNode(string databaseName, ObjectExplorerSession session)
         {
             Assert.NotNull(session);
             Assert.NotNull(session.Root);
@@ -147,7 +168,7 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.ObjectExplorer
             Assert.Equal(nodeInfo.IsLeaf, false);
             Assert.Equal(nodeInfo.NodeType, NodeTypes.Database.ToString());
             Assert.True(nodeInfo.Label.Contains(databaseName));
-            var children = session.Root.Expand();
+            var children = await _service.ExpandNode(session, session.Root.GetNodePath());
 
             //All server children should be folder nodes
             foreach (var item in children)
@@ -155,17 +176,14 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.ObjectExplorer
                 Assert.Equal(item.NodeType, "Folder");
             }
 
-            var tablesRoot = children.FirstOrDefault(x => x.NodeTypeId == NodeTypes.Tables);
+            var tablesRoot = children.FirstOrDefault(x => x.Label == SR.SchemaHierarchy_Tables);
             Assert.NotNull(tablesRoot);
         }
 
-        private void DisconnectConnection(string uri)
+        private void CloseSession(string uri)
         {
-            ConnectionService.Instance.Disconnect(new DisconnectParams
+            _service.CloseSession(uri);
-            {
+            Console.WriteLine($"Session closed uri:{uri}");
-                OwnerUri = uri,
-                Type = ConnectionType.Default
-            });
         }
 
         private async Task ExpandTree(NodeInfo node, ObjectExplorerSession session, StringBuilder stringBuilder = null, bool verifySystemObjects = false)
@@ -248,24 +266,22 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.ObjectExplorer
             }
         }
 
-        private async Task<bool> VerifyObjectExplorerTest(string databaseName, string queryFileName, string uri, string baselineFileName, bool verifySystemObjects = false)
+        private async Task<bool> VerifyObjectExplorerTest(string databaseName, string testDbPrefix, string queryFileName, string baselineFileName, bool verifySystemObjects = false)
         {
             var query = string.IsNullOrEmpty(queryFileName) ? string.Empty : LoadScript(queryFileName);
             StringBuilder stringBuilder = new StringBuilder();
-            SqlTestDb testDb = SqlTestDb.CreateNew(TestServerType.OnPrem, false, databaseName, query, uri);
+            await RunTest(databaseName, query, testDbPrefix, async (testDbName, session) =>
-            var session = await CreateSession(testDb.DatabaseName, uri);
+            {
-            await ExpandServerNodeAndVerifyDatabaseHierachy(testDb.DatabaseName, session, false);
+                await ExpandServerNodeAndVerifyDatabaseHierachy(testDbName, session, false);
                 await ExpandTree(session.Root.ToNodeInfo(), session, stringBuilder, verifySystemObjects);
-
                 string baseline = string.IsNullOrEmpty(baselineFileName) ? string.Empty : LoadBaseLine(baselineFileName);
                 if (!string.IsNullOrEmpty(baseline))
                 {
                     string actual = stringBuilder.ToString();
                     BaselinedTest.CompareActualWithBaseline(actual, baseline);
                 }
-            _service.CloseSession(session.Uri);
+            });
-            Thread.Sleep(3000);
+           
-            testDb.Cleanup();
             return true;
         }
 
@@ -316,6 +332,5 @@ namespace Microsoft.SqlTools.ServiceLayer.IntegrationTests.ObjectExplorer
             FileInfo inputFile = GetBaseLineFile(fileName);
             return TestUtilities.ReadTextAndNormalizeLineEndings(inputFile.FullName);
         }
-
     }
 }

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

@@ -9,6 +9,7 @@ using Microsoft.SqlTools.ServiceLayer.Connection;
 using Microsoft.SqlTools.ServiceLayer.Connection.Contracts;
 using Xunit;
 using System.Data.SqlClient;
+using System.Threading.Tasks;
 
 namespace Microsoft.SqlTools.ServiceLayer.Test.Common
 {
@@ -59,16 +60,29 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.Common
             string databaseName = null,
             string query = null,
             string dbNamePrefix = null)
+        {
+            return CreateNewAsync(serverType, doNotCleanupDb, databaseName, query, dbNamePrefix).Result;
+        }
+
+        /// <summary>
+        /// Create the test db if not already exists
+        /// </summary>
+        public static async Task<SqlTestDb> CreateNewAsync(
+            TestServerType serverType,
+            bool doNotCleanupDb = false,
+            string databaseName = null,
+            string query = null,
+            string dbNamePrefix = null)
         {
             SqlTestDb testDb = new SqlTestDb();
 
             databaseName = databaseName ?? GetUniqueDBName(dbNamePrefix);
             string createDatabaseQuery = Scripts.CreateDatabaseQuery.Replace("#DatabaseName#", databaseName);
-            TestServiceProvider.Instance.RunQuery(serverType, MasterDatabaseName, createDatabaseQuery);
+            await TestServiceProvider.Instance.RunQueryAsync(serverType, MasterDatabaseName, createDatabaseQuery);
             Console.WriteLine(string.Format(CultureInfo.InvariantCulture, "Test database '{0}' is created", databaseName));
             if (!string.IsNullOrEmpty(query))
             {
-                TestServiceProvider.Instance.RunQuery(serverType, databaseName, query);
+                await TestServiceProvider.Instance.RunQueryAsync(serverType, databaseName, query);
                 Console.WriteLine(string.Format(CultureInfo.InvariantCulture, "Test database '{0}' SQL types are created", databaseName));
             }
             testDb.DatabaseName = databaseName;
@@ -107,6 +121,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.Common
         }
 
         public void Cleanup()
+        {
+            CleanupAsync().Wait();
+        }
+
+        public async Task CleanupAsync()
         {
             try
             {
@@ -116,7 +135,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.Common
                         (ServerType == TestServerType.Azure ? Scripts.DropDatabaseIfExistAzure : Scripts.DropDatabaseIfExist), DatabaseName);
 
                     Console.WriteLine(string.Format(CultureInfo.InvariantCulture, "Cleaning up database {0}", DatabaseName));
-                    TestServiceProvider.Instance.RunQuery(ServerType, MasterDatabaseName, dropDatabaseQuery);
+                    await TestServiceProvider.Instance.RunQueryAsync(ServerType, MasterDatabaseName, dropDatabaseQuery);
                 }
             }
             catch (Exception ex)

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

@@ -79,6 +79,15 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.Common
         /// Runs a query by calling the services directly (not using the test driver) 
         /// </summary>
         public void RunQuery(TestServerType serverType, string databaseName, string queryText, bool throwOnError = false)
+        {
+            RunQueryAsync(serverType, databaseName, queryText, throwOnError).Wait();
+
+        }
+
+        /// <summary>
+        /// Runs a query by calling the services directly (not using the test driver) 
+        /// </summary>
+        public async Task RunQueryAsync(TestServerType serverType, string databaseName, string queryText, bool throwOnError = false)
         {
             string uri = "";
             using (SelfCleaningTempFile queryTempFile = new SelfCleaningTempFile())
@@ -88,7 +97,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.Common
                 ConnectionInfo connInfo = InitLiveConnectionInfo(serverType, databaseName, uri);
                 Query query = new Query(queryText, connInfo, new QueryExecutionSettings(), MemoryFileSystem.GetFileStreamFactory());
                 query.Execute();
-                query.ExecutionTask.Wait();
+                await query.ExecutionTask;
 
                 if (throwOnError)
                 {

test/Microsoft.SqlTools.ServiceLayer.UnitTests/ObjectExplorer/ObjectExplorerServiceTests.cs

@@ -4,6 +4,8 @@
 //
 
 using System;
+using System.Data.SqlClient;
+using System.Linq;
 using System.Threading.Tasks;
 using Microsoft.SqlTools.Hosting.Protocol;
 using Microsoft.SqlTools.ServiceLayer.Connection;
@@ -11,7 +13,6 @@ using Microsoft.SqlTools.ServiceLayer.Connection.Contracts;
 using Microsoft.SqlTools.ServiceLayer.ObjectExplorer;
 using Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Contracts;
 using Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes;
-using Microsoft.SqlTools.ServiceLayer.Test.Common;
 using Microsoft.SqlTools.ServiceLayer.UnitTests.Utility;
 using Moq;
 using Xunit;
@@ -116,6 +117,171 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ObjectExplorer
             await CreateSessionRequestAndVerifyServerNodeHelper(details);
         }
 
+        [Fact]
+        public async Task ExpandNodeGivenValidSessionShouldReturnTheNodeChildren()
+        {
+            await ExpandAndVerifyServerNodes();
+        }
+
+        [Fact]
+        public async Task RefreshNodeGivenValidSessionShouldReturnTheNodeChildren()
+        {
+            await RefreshAndVerifyServerNodes();
+        }
+
+        [Fact]
+        public async Task ExpandNodeGivenInvalidSessionShouldReturnEmptyList()
+        {
+            ExpandParams expandParams = new ExpandParams()
+            {
+                SessionId = "invalid session is",
+                NodePath = "Any path"
+            };
+
+            // when expanding
+            // then expect the nodes are server children 
+            await RunAndVerify<ExpandResponse>(
+                test: (requestContext) => service.HandleExpandRequest(expandParams, requestContext),
+                verify: (actual =>
+                {
+                    Assert.Equal(actual.SessionId, expandParams.SessionId);
+                    Assert.Null(actual.Nodes);
+                }));
+        }
+
+        [Fact]
+        public async Task RefreshNodeGivenInvalidSessionShouldReturnEmptyList()
+        {
+            RefreshParams expandParams = new RefreshParams()
+            {
+                SessionId = "invalid session is",
+                NodePath = "Any path"
+            };
+
+            // when expanding
+            // then expect the nodes are server children 
+            await RunAndVerify<RefreshResponse>(
+                test: (requestContext) => service.HandleRefreshRequest(expandParams, requestContext),
+                verify: (actual =>
+                {
+                    Assert.Equal(actual.SessionId, expandParams.SessionId);
+                    Assert.Null(actual.Nodes);
+                }));
+        }
+
+        [Fact]
+        public async Task CloseSessionGivenInvalidSessionShouldReturnEmptyList()
+        {
+            CloseSessionParams closeSessionParamsparams = new CloseSessionParams()
+            {
+                SessionId = "invalid session is",
+            };
+
+            // when expanding
+            // then expect the nodes are server children 
+            await RunAndVerify<CloseSessionResponse>(
+                test: (requestContext) => service.HandleCloseSessionRequest(closeSessionParamsparams, requestContext),
+                verify: (actual =>
+                {
+                    Assert.Equal(actual.SessionId, closeSessionParamsparams.SessionId);
+                    Assert.False(actual.Success);
+                }));
+        }
+
+        [Fact]
+        public async Task CloseSessionGivenValidSessionShouldCloseTheSessionAndDisconnect()
+        {
+            var session = await CreateSession();
+            CloseSessionParams closeSessionParamsparams = new CloseSessionParams()
+            {
+                SessionId = session.SessionId,
+            };
+
+            // when expanding
+            // then expect the nodes are server children 
+            await RunAndVerify<CloseSessionResponse>(
+                test: (requestContext) => service.HandleCloseSessionRequest(closeSessionParamsparams, requestContext),
+                verify: (actual =>
+                {
+                    Assert.Equal(actual.SessionId, closeSessionParamsparams.SessionId);
+                    Assert.True(actual.Success);
+                    Assert.False(service.SessionIds.Contains(session.SessionId));
+                }));
+
+            connectionServiceMock.Verify(c => c.Disconnect(It.IsAny<DisconnectParams>()));
+        }
+
+        private async Task<CreateSessionResponse> CreateSession()
+        {
+            ConnectionDetails details = new ConnectionDetails()
+            {
+                UserName = "user",
+                Password = "password",
+                DatabaseName = "msdb",
+                ServerName = "serverName"
+            };
+
+            serviceHostMock.AddEventHandling(ConnectionCompleteNotification.Type, null);
+            CreateSessionResponse result = default(CreateSessionResponse);
+            var contextMock = RequestContextMocks.Create<CreateSessionResponse>(r => result = r).AddErrorHandling(null);
+            connectionServiceMock.Setup(c => c.Connect(It.IsAny<ConnectParams>()))
+                .Returns((ConnectParams connectParams) => Task.FromResult(GetCompleteParamsForConnection(connectParams.OwnerUri, details)));
+
+            ConnectionInfo connectionInfo = new ConnectionInfo(null, null, null);
+            string fakeConnectionString = "Data Source=server;Initial Catalog=database;Integrated Security=False;User Id=user";
+            connectionInfo.AddConnection("Default", new SqlConnection(fakeConnectionString));
+            connectionServiceMock.Setup((c => c.TryFindConnection(It.IsAny<string>(), out connectionInfo))).
+                OutCallback((string t, out ConnectionInfo v) => v = connectionInfo)
+                .Returns(true);
+
+            connectionServiceMock.Setup(c => c.Disconnect(It.IsAny<DisconnectParams>())).Returns(true);
+            await service.HandleCreateSessionRequest(details, contextMock.Object);
+
+            return result;
+        }
+
+        private async Task ExpandAndVerifyServerNodes()
+        {
+            var session = await CreateSession();
+            ExpandParams expandParams = new ExpandParams()
+            {
+                SessionId = session.SessionId,
+                NodePath = session.RootNode.NodePath
+            };
+
+            // when expanding
+            // then expect the nodes are server children 
+            await RunAndVerify<ExpandResponse>(
+                test: (requestContext) => service.HandleExpandRequest(expandParams, requestContext),
+                verify: (actual =>
+                {
+                    Assert.Equal(actual.SessionId, session.SessionId);
+                    Assert.NotNull(actual.SessionId);
+                    VerifyServerNodeChildren(actual.Nodes);
+                }));
+        }
+
+        private async Task RefreshAndVerifyServerNodes()
+        {
+            var session = await CreateSession();
+            RefreshParams expandParams = new RefreshParams()
+            {
+                SessionId = session.SessionId,
+                NodePath = session.RootNode.NodePath
+            };
+
+            // when expanding
+            // then expect the nodes are server children 
+            await RunAndVerify<RefreshResponse>(
+                test: (requestContext) => service.HandleRefreshRequest(expandParams, requestContext),
+                verify: (actual =>
+                {
+                    Assert.Equal(actual.SessionId, session.SessionId);
+                    Assert.NotNull(actual.SessionId);
+                    VerifyServerNodeChildren(actual.Nodes);
+                }));
+        }
+
         private async Task CreateSessionRequestAndVerifyServerNodeHelper(ConnectionDetails details)
         {
             serviceHostMock.AddEventHandling(ConnectionCompleteNotification.Type, null);
@@ -150,6 +316,13 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ObjectExplorer
             Assert.False(serverNode.IsLeaf);
         }
 
+        private void VerifyServerNodeChildren(NodeInfo[] children)
+        {
+            Assert.NotNull(children);
+            Assert.True(children.Count() == 3);
+            Assert.True(children.All((x => x.NodeType == "Folder")));
+        }
+
         private static ConnectionCompleteParams GetCompleteParamsForConnection(string uri, ConnectionDetails details)
         {
             return new ConnectionCompleteParams()

test/Microsoft.SqlTools.ServiceLayer.UnitTests/ObjectExplorer/ObjectExplorerTestBase.cs

@@ -6,6 +6,8 @@
 using Microsoft.SqlTools.Extensibility;
 using Microsoft.SqlTools.ServiceLayer.Connection;
 using Microsoft.SqlTools.ServiceLayer.ObjectExplorer;
+using Microsoft.SqlTools.ServiceLayer.ObjectExplorer.Nodes;
+using Microsoft.SqlTools.ServiceLayer.ObjectExplorer.SmoModel;
 
 namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ObjectExplorer
 {
@@ -35,7 +37,8 @@ namespace Microsoft.SqlTools.ServiceLayer.UnitTests.ObjectExplorer
         {
             CreateProvider()
                 .RegisterSingleService(connService)
-                .RegisterSingleService(new ObjectExplorerService());
+                .RegisterSingleService(new ObjectExplorerService())
+                .RegisterSingleService<ChildFactory>(new ServerChildFactory());
 
             // Create the service using the service provider, which will initialize dependencies
             return ServiceProvider.GetService<ObjectExplorerService>();