Merge branch 'dev' into bug/addGetSignatureHelpTest

test/Microsoft.SqlTools.ServiceLayer.Test/LanguageServer/AutocompleteTests.cs

@@ -39,8 +39,12 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
 
         private Mock<RequestContext<CompletionItem[]>> requestContext;
 
+        private Mock<ScriptFile> scriptFile;
+
         private Mock<IBinder> binder; 
 
+        private ScriptParseInfo scriptParseInfo;  
+
         private TextDocumentPosition textDocument;
 
         private void InitializeTestObjects()
@@ -60,14 +64,14 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
             WorkspaceService<SqlToolsSettings>.Instance.CurrentSettings = new SqlToolsSettings();
 
             // set up file for returning the query
-            var fileMock = new Mock<ScriptFile>();
-            fileMock.SetupGet(file => file.Contents).Returns(Common.StandardQuery);
-            fileMock.SetupGet(file => file.ClientFilePath).Returns(this.testScriptUri);
+            scriptFile = new Mock<ScriptFile>();
+            scriptFile.SetupGet(file => file.Contents).Returns(Common.StandardQuery);
+            scriptFile.SetupGet(file => file.ClientFilePath).Returns(this.testScriptUri);
 
             // set up workspace mock
             workspaceService = new Mock<WorkspaceService<SqlToolsSettings>>();
             workspaceService.Setup(service => service.Workspace.GetFile(It.IsAny<string>()))
-                .Returns(fileMock.Object);
+                .Returns(scriptFile.Object);
         
             // setup binding queue mock
             bindingQueue = new Mock<ConnectedBindingQueue>();
@@ -93,16 +97,113 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
                 It.IsAny<string>(),
                 It.IsAny<BindMode>()));
             
-            var testScriptParseInfo = new ScriptParseInfo();        
-            LanguageService.Instance.AddOrUpdateScriptParseInfo(this.testScriptUri, testScriptParseInfo);      
-            testScriptParseInfo.IsConnected = true;            
-            testScriptParseInfo.ConnectionKey = LanguageService.Instance.BindingQueue.AddConnectionContext(connectionInfo);
+            scriptParseInfo = new ScriptParseInfo();        
+            LanguageService.Instance.AddOrUpdateScriptParseInfo(this.testScriptUri, scriptParseInfo);      
+            scriptParseInfo.IsConnected = true;            
+            scriptParseInfo.ConnectionKey = LanguageService.Instance.BindingQueue.AddConnectionContext(connectionInfo);
 
             // setup the binding context object
             ConnectedBindingContext bindingContext = new ConnectedBindingContext();
             bindingContext.Binder = binder.Object;
             bindingContext.MetadataDisplayInfoProvider = new MetadataDisplayInfoProvider();
-            LanguageService.Instance.BindingQueue.BindingContextMap.Add(testScriptParseInfo.ConnectionKey, bindingContext);                
+            LanguageService.Instance.BindingQueue.BindingContextMap.Add(scriptParseInfo.ConnectionKey, bindingContext);                
+        }
+
+        [Fact]
+        public void HandleCompletionRequestDisabled()
+        {
+            InitializeTestObjects();
+            WorkspaceService<SqlToolsSettings>.Instance.CurrentSettings.SqlTools.IntelliSense.EnableIntellisense = false;            
+            Assert.NotNull(LanguageService.HandleCompletionRequest(null, null));
+        }
+
+        [Fact]
+        public void HandleCompletionResolveRequestDisabled()
+        {
+            InitializeTestObjects();
+            WorkspaceService<SqlToolsSettings>.Instance.CurrentSettings.SqlTools.IntelliSense.EnableIntellisense = false;            
+            Assert.NotNull(LanguageService.HandleCompletionResolveRequest(null, null));
+        }
+
+        [Fact]
+        public void HandleSignatureHelpRequestDisabled()
+        {
+            InitializeTestObjects();
+            WorkspaceService<SqlToolsSettings>.Instance.CurrentSettings.SqlTools.IntelliSense.EnableIntellisense = false;            
+            Assert.NotNull(LanguageService.HandleSignatureHelpRequest(null, null));
+        }               
+
+        [Fact]
+        public void AddOrUpdateScriptParseInfoNullUri()
+        {
+            InitializeTestObjects();
+            LanguageService.Instance.AddOrUpdateScriptParseInfo("abracadabra", scriptParseInfo);
+            Assert.True(LanguageService.Instance.ScriptParseInfoMap.ContainsKey("abracadabra"));
+        }
+
+        [Fact]
+        public void GetDefinitionInvalidTextDocument()
+        {
+            InitializeTestObjects();
+            textDocument.TextDocument.Uri = "invaliduri";
+            Assert.Null(LanguageService.Instance.GetDefinition(textDocument, null, null));
+        }
+
+        [Fact]
+        public void RemoveScriptParseInfoNullUri()
+        {
+            InitializeTestObjects();
+            Assert.False(LanguageService.Instance.RemoveScriptParseInfo("abc123"));
+        }
+
+        [Fact]
+        public void IsPreviewWindowNullScriptFileTest()
+        {
+            InitializeTestObjects();
+            Assert.False(LanguageService.Instance.IsPreviewWindow(null));
+        }
+
+        [Fact]
+        public void GetCompletionItemsInvalidTextDocument()
+        {
+            InitializeTestObjects();
+            textDocument.TextDocument.Uri = "somethinggoeshere";
+            Assert.True(LanguageService.Instance.GetCompletionItems(textDocument, scriptFile.Object, null).Length > 0);
+        }
+
+        [Fact]
+        public void GetDiagnosticFromMarkerTest()
+        {
+            var scriptFileMarker = new ScriptFileMarker()
+            {
+                Message = "Message",
+                Level = ScriptFileMarkerLevel.Error,
+                ScriptRegion = new ScriptRegion()
+                {
+                    File = "file://nofile.sql",
+                    StartLineNumber = 1,
+                    StartColumnNumber = 1,
+                    StartOffset = 0,
+                    EndLineNumber = 1,
+                    EndColumnNumber = 1,
+                    EndOffset = 0
+                }
+            }; 
+            var diagnostic = DiagnosticsHelper.GetDiagnosticFromMarker(scriptFileMarker);
+            Assert.Equal(diagnostic.Message, scriptFileMarker.Message);
+        }
+
+        [Fact]
+        public void MapDiagnosticSeverityTest()
+        {
+            var level = ScriptFileMarkerLevel.Error;
+            Assert.Equal(DiagnosticsHelper.MapDiagnosticSeverity(level), DiagnosticSeverity.Error);
+            level = ScriptFileMarkerLevel.Warning;
+            Assert.Equal(DiagnosticsHelper.MapDiagnosticSeverity(level), DiagnosticSeverity.Warning);
+            level = ScriptFileMarkerLevel.Information;
+            Assert.Equal(DiagnosticsHelper.MapDiagnosticSeverity(level), DiagnosticSeverity.Information);
+            level = (ScriptFileMarkerLevel)100;
+            Assert.Equal(DiagnosticsHelper.MapDiagnosticSeverity(level), DiagnosticSeverity.Error);
         }
 
         /// <summary>

test/Microsoft.SqlTools.ServiceLayer.Test/LanguageServer/LanguageServiceTests.cs

@@ -3,8 +3,10 @@
 // Licensed under the MIT license. See LICENSE file in the project root for full license information.
 //
 
+using Microsoft.SqlServer.Management.SqlParser.Parser;
 using Microsoft.SqlTools.ServiceLayer.Connection;
 using Microsoft.SqlTools.ServiceLayer.LanguageServices;
+using Microsoft.SqlTools.ServiceLayer.LanguageServices.Completion;
 using Microsoft.SqlTools.ServiceLayer.LanguageServices.Contracts;
 using Microsoft.SqlTools.ServiceLayer.Workspace.Contracts;
 using Microsoft.SqlTools.Test.Utility;
@@ -147,6 +149,56 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServer
             Assert.Null(signatureHelp);
         }
 
+        [Fact]
+        public void EmptyCompletionListTest()
+        {           
+            Assert.Equal(AutoCompleteHelper.EmptyCompletionList.Length, 0);
+        }
+
+        [Fact]
+        public void SetWorkspaceServiceInstanceTest()
+        {           
+            AutoCompleteHelper.WorkspaceServiceInstance = null;
+            // workspace will be recreated if it's set to null
+            Assert.NotNull(AutoCompleteHelper.WorkspaceServiceInstance);
+        }
+
+        internal class TestScriptDocumentInfo : ScriptDocumentInfo
+        {
+            public TestScriptDocumentInfo(TextDocumentPosition textDocumentPosition, ScriptFile scriptFile, ScriptParseInfo scriptParseInfo)
+                :base(textDocumentPosition, scriptFile, scriptParseInfo)
+            {
+                
+            }
+            
+            public override string TokenText
+            {
+                get
+                {
+                    return "doesntmatchanythingintheintellisensedefaultlist";
+                }
+            }
+        }
+
+        [Fact]
+        public void GetDefaultCompletionListWithNoMatchesTest()
+        {           
+            var scriptFile = new ScriptFile();
+            scriptFile.SetFileContents("koko wants a bananas");
+
+            ScriptParseInfo scriptInfo = new ScriptParseInfo { IsConnected = false };
+
+            var scriptDocumentInfo = new TestScriptDocumentInfo(
+                new TextDocumentPosition()
+                {
+                    TextDocument = new TextDocumentIdentifier() {  Uri = TestObjects.ScriptUri  },
+                    Position = new Position() { Line = 0, Character = 0 }
+                }, scriptFile, scriptInfo);
+      
+            AutoCompleteHelper.GetDefaultCompletionItems(scriptDocumentInfo, false);
+
+        }
+
         #endregion
 
         #region "General Language Service tests"
@@ -170,11 +222,10 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServer
 
             connInfo = TestObjects.InitLiveConnectionInfo(out scriptFile);
         }
-
+        
         /// <summary>
         /// Test the service initialization code path and verify nothing throws
         /// </summary>
-        // Test is causing failures in build lab..investigating to reenable
         [Fact]
         public void ServiceInitialization()
         {
@@ -190,12 +241,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServer
             Assert.True(LanguageService.ConnectionServiceInstance != null);
             Assert.True(LanguageService.Instance.CurrentSettings != null);
             Assert.True(LanguageService.Instance.CurrentWorkspace != null);
-        }        
+        }  
 
         /// <summary>
         /// Test the service initialization code path and verify nothing throws
         /// </summary>
-        // Test is causing failures in build lab..investigating to reenable
         [Fact]
         public void PrepopulateCommonMetadata()
         {

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

@@ -5,22 +5,27 @@
 using System;
 using System.IO;
 using System.Collections.Generic;
+using System.Threading;
 using System.Threading.Tasks;
 using System.Runtime.InteropServices;
+using Microsoft.SqlServer.Management.Common;
 using Microsoft.SqlServer.Management.SqlParser.Binder;
+using Microsoft.SqlServer.Management.SqlParser.Intellisense;
 using Microsoft.SqlServer.Management.SqlParser.MetadataProvider;
 using Microsoft.SqlServer.Management.SqlParser.Parser;
 using Microsoft.SqlTools.ServiceLayer.Connection;
 using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol;
+using Microsoft.SqlTools.ServiceLayer.Hosting.Protocol.Contracts;
 using Microsoft.SqlTools.ServiceLayer.LanguageServices;
+using Microsoft.SqlTools.ServiceLayer.LanguageServices.Contracts;
 using Microsoft.SqlTools.ServiceLayer.SqlContext;
 using Microsoft.SqlTools.ServiceLayer.Test.QueryExecution;
 using Microsoft.SqlTools.ServiceLayer.Workspace;
 using Microsoft.SqlTools.ServiceLayer.Workspace.Contracts;
 using Microsoft.SqlTools.Test.Utility;
-using Location = Microsoft.SqlTools.ServiceLayer.Workspace.Contracts.Location;
 using Moq;
 using Xunit;
+using Location = Microsoft.SqlTools.ServiceLayer.Workspace.Contracts.Location;
 
 namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
 {
@@ -89,6 +94,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
             requestContext = new Mock<RequestContext<Location[]>>();
             requestContext.Setup(rc => rc.SendResult(It.IsAny<Location[]>()))
                 .Returns(Task.FromResult(0));
+            requestContext.Setup(rc => rc.SendError(It.IsAny<DefinitionError>())).Returns(Task.FromResult(0));;
+            requestContext.Setup(r => r.SendEvent(It.IsAny<EventType<TelemetryParams>>(), It.IsAny<TelemetryParams>())).Returns(Task.FromResult(0));;
+            requestContext.Setup(r => r.SendEvent(It.IsAny<EventType<StatusChangeParams>>(), It.IsAny<StatusChangeParams>())).Returns(Task.FromResult(0));;
 
             // setup the IBinder mock
             binder = new Mock<IBinder>();
@@ -99,7 +107,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
 
             var testScriptParseInfo = new ScriptParseInfo();
             LanguageService.Instance.AddOrUpdateScriptParseInfo(this.testScriptUri, testScriptParseInfo);
-            testScriptParseInfo.IsConnected = true;
+            testScriptParseInfo.IsConnected = false;
             testScriptParseInfo.ConnectionKey = LanguageService.Instance.BindingQueue.AddConnectionContext(connectionInfo);
 
             // setup the binding context object
@@ -111,18 +119,19 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
 
 
         /// <summary>
-        /// Tests the definition event handler. When called with no active connection, no definition is sent
+        /// Tests the definition event handler. When called with no active connection, an error is sent
         /// </summary>
         [Fact]
-        public void DefinitionsHandlerWithNoConnectionTest()
+        public async Task DefinitionsHandlerWithNoConnectionTest()
         {
+            TestObjects.InitializeTestServices();
             InitializeTestObjects();
-            // request the completion list
-            Task handleCompletion = LanguageService.HandleDefinitionRequest(textDocument, requestContext.Object);
-            handleCompletion.Wait(TaskTimeout);
-
-            // verify that send result was not called
+            // request definition
+            var definitionTask = await Task.WhenAny(LanguageService.HandleDefinitionRequest(textDocument, requestContext.Object), Task.Delay(TaskTimeout));
+            await definitionTask;
+            // verify that send result was not called and send error was called
             requestContext.Verify(m => m.SendResult(It.IsAny<Location[]>()), Times.Never());
+            requestContext.Verify(m => m.SendError(It.IsAny<DefinitionError>()), Times.Once());
         }
 
         /// <summary>
@@ -132,7 +141,7 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
         public void GetLocationFromFileForValidFilePathTest()
         {
             String filePath = RuntimeInformation.IsOSPlatform(OSPlatform.Windows) ? "C:\\test\\script.sql" : "/test/script.sql";
-            PeekDefinition peekDefinition = new PeekDefinition(null);
+            PeekDefinition peekDefinition = new PeekDefinition(null, null);
             Location[] locations = peekDefinition.GetLocationFromFile(filePath, 0);
 
             String expectedFilePath = RuntimeInformation.IsOSPlatform(OSPlatform.Windows) ? "file:///C:/test/script.sql" : "file:/test/script.sql";
@@ -145,11 +154,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
         [Fact]
         public void GetSchemaFromDatabaseQualifiedNameWithValidNameTest()
         {
-            PeekDefinition peekDefinition = new PeekDefinition(null);
+            PeekDefinition peekDefinition = new PeekDefinition(null, null);
             string validDatabaseQualifiedName = "master.test.test_table";
             string objectName = "test_table";
             string expectedSchemaName = "test";
-        
+
             string actualSchemaName = peekDefinition.GetSchemaFromDatabaseQualifiedName(validDatabaseQualifiedName, objectName);
             Assert.Equal(actualSchemaName, expectedSchemaName);
         }
@@ -161,11 +170,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
         [Fact]
         public void GetSchemaFromDatabaseQualifiedNameWithNoSchemaTest()
         {
-            PeekDefinition peekDefinition = new PeekDefinition(null);
+            PeekDefinition peekDefinition = new PeekDefinition(null, null);
             string validDatabaseQualifiedName = "test_table";
             string objectName = "test_table";
             string expectedSchemaName = "dbo";
-        
+
             string actualSchemaName = peekDefinition.GetSchemaFromDatabaseQualifiedName(validDatabaseQualifiedName, objectName);
             Assert.Equal(actualSchemaName, expectedSchemaName);
         }
@@ -176,15 +185,15 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
         [Fact]
         public void GetSchemaFromDatabaseQualifiedNameWithInvalidNameTest()
         {
-            PeekDefinition peekDefinition = new PeekDefinition(null);
+            PeekDefinition peekDefinition = new PeekDefinition(null, null);
             string validDatabaseQualifiedName = "x.y.z";
             string objectName = "test_table";
             string expectedSchemaName = "dbo";
-        
+
             string actualSchemaName = peekDefinition.GetSchemaFromDatabaseQualifiedName(validDatabaseQualifiedName, objectName);
             Assert.Equal(actualSchemaName, expectedSchemaName);
         }
-         
+
 #if LIVE_CONNECTION_TESTS
         /// <summary>
         /// Test get definition for a table object with active connection
@@ -192,10 +201,13 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
         [Fact]
         public void GetValidTableDefinitionTest()
         {
-            // Get live connectionInfo
+            // Get live connectionInfo and serverConnection
             ConnectionInfo connInfo = TestObjects.InitLiveConnectionInfoForDefinition();
-            PeekDefinition peekDefinition = new PeekDefinition(connInfo);
-            string objectName = "test_table";
+            ServerConnection serverConnection = TestObjects.InitLiveServerConnectionForDefinition(connInfo);
+
+            PeekDefinition peekDefinition = new PeekDefinition(serverConnection, connInfo);
+            string objectName = "spt_monitor";
+
             string schemaName = null;
             string objectType = "TABLE";
 
@@ -211,9 +223,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
         [Fact]
         public void GetTableDefinitionInvalidObjectTest()
         {
-            // Get live connectionInfo
+            // Get live connectionInfo and serverConnection
             ConnectionInfo connInfo = TestObjects.InitLiveConnectionInfoForDefinition();
-            PeekDefinition peekDefinition = new PeekDefinition(connInfo);
+            ServerConnection serverConnection = TestObjects.InitLiveServerConnectionForDefinition(connInfo);
+
+            PeekDefinition peekDefinition = new PeekDefinition(serverConnection, connInfo);
             string objectName = "test_invalid";
             string schemaName = null;
             string objectType = "TABLE";
@@ -229,25 +243,102 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
         [Fact]
         public void GetTableDefinitionWithSchemaTest()
         {
-            // Get live connectionInfo
+            // Get live connectionInfo and serverConnection
             ConnectionInfo connInfo = TestObjects.InitLiveConnectionInfoForDefinition();
-            PeekDefinition peekDefinition = new PeekDefinition(connInfo);
-            string objectName = "test_table";
+            ServerConnection serverConnection = TestObjects.InitLiveServerConnectionForDefinition(connInfo);
+
+            PeekDefinition peekDefinition = new PeekDefinition(serverConnection, connInfo);
+            string objectName = "spt_monitor";
+
             string schemaName = "dbo";
             string objectType = "TABLE";
 
             // Get locations for valid table object with schema name
             Location[] locations = peekDefinition.GetSqlObjectDefinition(peekDefinition.GetTableScripts, objectName, schemaName, objectType);
             Assert.NotNull(locations);
-            Cleanup(locations);            
+            Cleanup(locations);
         }
 
         /// <summary>
-        /// Test GetDefinition with an unsupported type(function)
+        /// Test GetDefinition with an unsupported type(schema - dbo). Expect a error result.
         /// </summary>
         [Fact]
-        public void GetUnsupportedDefinitionForFullScript()
+        public void GetUnsupportedDefinitionErrorTest()
         {
+            ScriptFile scriptFile;
+            TextDocumentPosition textDocument = new TextDocumentPosition
+            {
+                TextDocument = new TextDocumentIdentifier { Uri = OwnerUri },
+                Position = new Position
+                {
+                    Line = 0,
+                    // test for 'dbo'
+                    Character = 16
+                }
+            };
+            ConnectionInfo connInfo = TestObjects.InitLiveConnectionInfo(out scriptFile);
+            scriptFile.Contents = "select * from dbo.func ()";
+            var languageService = new LanguageService();
+            ScriptParseInfo scriptInfo = new ScriptParseInfo { IsConnected = true };
+            languageService.ScriptParseInfoMap.Add(OwnerUri, scriptInfo);
+
+            // When I call the language service
+            var result = languageService.GetDefinition(textDocument, scriptFile, connInfo);
+
+            // Then I expect null locations and an error to be reported
+            Assert.NotNull(result);
+            Assert.True(result.IsErrorResult);
+        }
+
+        /// <summary>
+        /// Get Definition for a object with no definition. Expect a error result
+        /// </summary>
+        [Fact]
+        public void GetDefinitionWithNoResultsFoundError()
+        {
+            ConnectionInfo connInfo = TestObjects.InitLiveConnectionInfoForDefinition();
+            ServerConnection serverConnection = TestObjects.InitLiveServerConnectionForDefinition(connInfo);
+
+            PeekDefinition peekDefinition = new PeekDefinition(serverConnection, connInfo);
+            string objectName = "from";
+
+            List<Declaration> declarations = new List<Declaration>();
+            DefinitionResult result = peekDefinition.GetScript(declarations, objectName, null);
+
+            Assert.NotNull(result);
+            Assert.True(result.IsErrorResult);
+            Assert.Equal(SR.PeekDefinitionNoResultsError, result.Message);
+        }
+    
+        /// <summary>
+        /// Test GetDefinition with a forced timeout. Expect a error result.
+        /// </summary>
+        [Fact]
+        public void GetDefinitionTimeoutTest()
+        {
+            // Given a binding queue that will automatically time out
+            var languageService = new LanguageService();
+            Mock<ConnectedBindingQueue> queueMock = new Mock<ConnectedBindingQueue>();
+            languageService.BindingQueue = queueMock.Object;
+            ManualResetEvent mre = new ManualResetEvent(true); // Do not block
+            Mock<QueueItem> itemMock = new Mock<QueueItem>();
+            itemMock.Setup(i => i.ItemProcessed).Returns(mre);
+            
+            DefinitionResult timeoutResult = null;
+            
+            queueMock.Setup(q => q.QueueBindingOperation(
+                It.IsAny<string>(),
+                It.IsAny<Func<IBindingContext, CancellationToken, object>>(),
+                It.IsAny<Func<IBindingContext, object>>(),
+                It.IsAny<int?>(),
+                It.IsAny<int?>()))
+            .Callback<string, Func<IBindingContext, CancellationToken, object>, Func<IBindingContext, object>, int?, int?>(
+                (key, bindOperation, timeoutOperation, blah, blah2) => 
+            {
+                timeoutResult = (DefinitionResult) timeoutOperation((IBindingContext)null);
+                itemMock.Object.Result = timeoutResult;
+            })
+            .Returns(() => itemMock.Object);
 
             ScriptFile scriptFile;
             TextDocumentPosition textDocument = new TextDocumentPosition
@@ -262,12 +353,17 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
             ConnectionInfo connInfo = TestObjects.InitLiveConnectionInfo(out scriptFile);
             scriptFile.Contents = "select * from dbo.func ()";
 
-            var languageService = LanguageService.Instance;
             ScriptParseInfo scriptInfo = new ScriptParseInfo { IsConnected = true };
             languageService.ScriptParseInfoMap.Add(OwnerUri, scriptInfo);
 
-            var locations = languageService.GetDefinition(textDocument, scriptFile, connInfo);
-            Assert.Null(locations);
+            // When I call the language service
+            var result = languageService.GetDefinition(textDocument, scriptFile, connInfo);
+
+            // Then I expect null locations and an error to be reported
+            Assert.NotNull(result);
+            Assert.True(result.IsErrorResult);
+            // Check timeout message
+            Assert.Equal(SR.PeekDefinitionTimedoutError, result.Message);
         }
 
         /// <summary>
@@ -277,7 +373,9 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
         public void GetValidViewDefinitionTest()
         {
             ConnectionInfo connInfo = TestObjects.InitLiveConnectionInfoForDefinition();
-            PeekDefinition peekDefinition = new PeekDefinition(connInfo);
+            ServerConnection serverConnection = TestObjects.InitLiveServerConnectionForDefinition(connInfo);
+
+            PeekDefinition peekDefinition = new PeekDefinition(serverConnection, connInfo);
             string objectName = "objects";
             string schemaName = "sys";
             string objectType = "VIEW";
@@ -293,8 +391,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
         [Fact]
         public void GetViewDefinitionInvalidObjectTest()
         {
+            // Get live connectionInfo and serverConnection
             ConnectionInfo connInfo = TestObjects.InitLiveConnectionInfoForDefinition();
-            PeekDefinition peekDefinition = new PeekDefinition(connInfo);
+            ServerConnection serverConnection = TestObjects.InitLiveServerConnectionForDefinition(connInfo);
+
+            PeekDefinition peekDefinition = new PeekDefinition(serverConnection, connInfo);
             string objectName = "objects";
             string schemaName = null;
             string objectType = "VIEW";
@@ -309,9 +410,13 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
         [Fact]
         public void GetStoredProcedureDefinitionTest()
         {
+            // Get live connectionInfo and serverConnection
             ConnectionInfo connInfo = TestObjects.InitLiveConnectionInfoForDefinition();
-            PeekDefinition peekDefinition = new PeekDefinition(connInfo);
-            string objectName = "SP1";
+            ServerConnection serverConnection = TestObjects.InitLiveServerConnectionForDefinition(connInfo);
+
+            PeekDefinition peekDefinition = new PeekDefinition(serverConnection, connInfo);
+            string objectName = "sp_MSrepl_startup";
+
             string schemaName = "dbo";
             string objectType = "PROCEDURE";
 
@@ -326,8 +431,11 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
         [Fact]
         public void GetStoredProcedureDefinitionFailureTest()
         {
+            // Get live connectionInfo and serverConnection
             ConnectionInfo connInfo = TestObjects.InitLiveConnectionInfoForDefinition();
-            PeekDefinition peekDefinition = new PeekDefinition(connInfo);
+            ServerConnection serverConnection = TestObjects.InitLiveServerConnectionForDefinition(connInfo);
+
+            PeekDefinition peekDefinition = new PeekDefinition(serverConnection, connInfo);
             string objectName = "SP2";
             string schemaName = "dbo";
             string objectType = "PROCEDURE";
@@ -342,9 +450,12 @@ namespace Microsoft.SqlTools.ServiceLayer.Test.LanguageServices
         [Fact]
         public void GetStoredProcedureDefinitionWithoutSchemaTest()
         {
+            // Get live connectionInfo and serverConnection
             ConnectionInfo connInfo = TestObjects.InitLiveConnectionInfoForDefinition();
-            PeekDefinition peekDefinition = new PeekDefinition(connInfo);
-            string objectName = "SP1";
+            ServerConnection serverConnection = TestObjects.InitLiveServerConnectionForDefinition(connInfo);
+
+            PeekDefinition peekDefinition = new PeekDefinition(serverConnection, connInfo);
+            string objectName = "sp_MSrepl_startup";
             string schemaName = null;
             string objectType = "PROCEDURE";