Fix crash and hang issues with Pester CodeLens feature

src/PowerShellEditorServices.Host/CodeLens/PesterCodeLensProvider.cs

@@ -27,18 +27,9 @@ public PesterCodeLensProvider(EditorSession editorSession)
         }
 
         private IEnumerable<CodeLens> GetPesterLens(
-            SymbolReference symbol,
+            PesterSymbolReference pesterSymbol,
             ScriptFile scriptFile)
         {
-            // Trim the Describe "" { from the symbol name
-            int startQuoteIndex = symbol.SourceLine.IndexOfAny(QuoteChars) + 1;
-            int endQuoteIndex =  symbol.SourceLine.LastIndexOfAny(QuoteChars);
-
-            string describeBlockName =
-                symbol.SourceLine.Substring(
-                    startQuoteIndex,
-                    endQuoteIndex - startQuoteIndex);
-
             var clientCommands = new ClientCommand[]
             {
                 new ClientCommand(
@@ -48,7 +39,7 @@ private IEnumerable<CodeLens> GetPesterLens(
                     {
                         scriptFile.ClientFilePath,
                         false, // Don't debug
-                        describeBlockName,
+                        pesterSymbol.TestName,
                     }),
 
                 new ClientCommand(
@@ -58,7 +49,7 @@ private IEnumerable<CodeLens> GetPesterLens(
                     {
                         scriptFile.ClientFilePath,
                         true, // Run in debugger
-                        describeBlockName,
+                        pesterSymbol.TestName,
                     }),
             };
 
@@ -68,7 +59,7 @@ private IEnumerable<CodeLens> GetPesterLens(
                         new CodeLens(
                             this,
                             scriptFile,
-                            symbol.ScriptRegion,
+                            pesterSymbol.ScriptRegion,
                             command));
         }
 
@@ -80,8 +71,10 @@ public CodeLens[] ProvideCodeLenses(ScriptFile scriptFile)
 
             var lenses =
                 symbols
-                    .Where(s => s.SymbolName.StartsWith("Describe"))
+                    .OfType<PesterSymbolReference>()
+                    .Where(s => s.Command == PesterCommandType.Describe)
                     .SelectMany(s => this.GetPesterLens(s, scriptFile))
+                    .Where(codeLens => codeLens != null)
                     .ToArray();
 
             return lenses;

src/PowerShellEditorServices/Symbols/PesterDocumentSymbolProvider.cs

@@ -16,7 +16,6 @@ namespace Microsoft.PowerShell.EditorServices.Symbols
     /// </summary>
     public class PesterDocumentSymbolProvider : FeatureProviderBase, IDocumentSymbolProvider
     {
-        private static char[] DefinitionTrimChars = new char[] { ' ', '{' };
 
         IEnumerable<SymbolReference> IDocumentSymbolProvider.ProvideDocumentSymbols(
             ScriptFile scriptFile)
@@ -30,33 +29,135 @@ IEnumerable<SymbolReference> IDocumentSymbolProvider.ProvideDocumentSymbols(
 
             var commandAsts = scriptFile.ScriptAst.FindAll(ast =>
             {
-                switch ((ast as CommandAst)?.GetCommandName()?.ToLower())
-                {
-                    case "describe":
-                    case "context":
-                    case "it":
-                        return true;
-
-                    default:
-                        return false;
-                }
+            CommandAst commandAst = ast as CommandAst;
+
+                return
+                    commandAst != null &&
+                    PesterSymbolReference.GetCommandType(commandAst.GetCommandName()).HasValue &&
+                    commandAst.CommandElements.Count >= 2;
             },
             true);
 
             return commandAsts.Select(
-                ast => {
-                    var testDefinitionLine =
+                ast =>
+                {
+                    // By this point we know the Ast is a CommandAst with 2 or more CommandElements
+                    int testNameParamIndex = 1;
+                    CommandAst testAst = (CommandAst)ast;
+
+                    // The -Name parameter
+                    for (int i = 1; i < testAst.CommandElements.Count; i++)
+                    {
+                        CommandParameterAst paramAst = testAst.CommandElements[i] as CommandParameterAst;
+                        if (paramAst != null &&
+                            paramAst.ParameterName.Equals("Name", StringComparison.OrdinalIgnoreCase))
+                        {
+                            testNameParamIndex = i + 1;
+                            break;
+                        }
+                    }
+
+                    if (testNameParamIndex > testAst.CommandElements.Count - 1)
+                    {
+                        return null;
+                    }
+
+                    StringConstantExpressionAst stringAst =
+                        testAst.CommandElements[testNameParamIndex] as StringConstantExpressionAst;
+
+                    if (stringAst == null)
+                    {
+                        return null;
+                    }
+
+                    string testDefinitionLine =
                         scriptFile.GetLine(
                             ast.Extent.StartLineNumber);
 
                     return
-                        new SymbolReference(
-                            SymbolType.Function,
-                            testDefinitionLine.TrimEnd(DefinitionTrimChars),
-                            ast.Extent,
-                            scriptFile.FilePath,
-                            testDefinitionLine);
-                });
+                        new PesterSymbolReference(
+                            scriptFile,
+                            testAst.GetCommandName(),
+                            testDefinitionLine,
+                            stringAst.Value,
+                            ast.Extent);
+
+                }).Where(s => s != null);
+        }
+    }
+
+    /// <summary>
+    /// Defines command types for Pester test blocks.
+    /// </summary>
+    public enum PesterCommandType
+    {
+        /// <summary>
+        /// Identifies a Describe block.
+        /// </summary>
+        Describe,
+
+        /// <summary>
+        /// Identifies a Context block.
+        /// </summary>
+        Context,
+
+        /// <summary>
+        /// Identifies an It block.
+        /// </summary>
+        It
+    }
+
+    /// <summary>
+    /// Provides a specialization of SymbolReference containing
+    /// extra information about Pester test symbols.
+    /// </summary>
+    public class PesterSymbolReference : SymbolReference
+    {
+        private static char[] DefinitionTrimChars = new char[] { ' ', '{' };
+
+        /// <summary>
+        /// Gets the name of the test
+        /// </summary>
+        public string TestName { get; private set; }
+
+        /// <summary>
+        /// Gets the test's command type.
+        /// </summary>
+        public PesterCommandType Command { get; private set; }
+
+        internal PesterSymbolReference(
+            ScriptFile scriptFile,
+            string commandName,
+            string testLine,
+            string testName,
+            IScriptExtent scriptExtent)
+                : base(
+                    SymbolType.Function,
+                    testLine.TrimEnd(DefinitionTrimChars),
+                    scriptExtent,
+                    scriptFile.FilePath,
+                    testLine)
+        {
+            this.Command = GetCommandType(commandName).Value;
+            this.TestName = testName;
+        }
+
+        internal static PesterCommandType? GetCommandType(string commandName)
+        {
+            switch (commandName.ToLower())
+            {
+                case "describe":
+                    return PesterCommandType.Describe;
+
+                case "context":
+                    return PesterCommandType.Context;
+
+                case "it":
+                    return PesterCommandType.It;
+
+                default:
+                    return null;
+            }
         }
     }
 }