Refactor pester script detection

src/PowerShellEditorServices/Symbols/PesterDocumentSymbolProvider.cs

@@ -27,63 +27,104 @@ IEnumerable<SymbolReference> IDocumentSymbolProvider.ProvideDocumentSymbols(
                 return Enumerable.Empty<SymbolReference>();
             }
 
-            var commandAsts = scriptFile.ScriptAst.FindAll(ast =>
-            {
+            // Find plausible Pester commands
+            IEnumerable<Ast> commandAsts = scriptFile.ScriptAst.FindAll(IsNamedCommandWithArguments, true);
+
+            return commandAsts.OfType<CommandAst>()
+                              .Where(IsPesterCommand)
+                              .Select(ast => ConvertPesterAstToSymbolReference(scriptFile, ast))
+                              .Where(pesterSymbol => pesterSymbol?.TestName != null);
+        }
+
+        /// <summary>
+        /// Test if the given Ast is a regular CommandAst with arguments
+        /// </summary>
+        /// <param name="ast">the PowerShell Ast to test</param>
+        /// <returns>true if the Ast represents a PowerShell command with arguments, false otherwise</returns>
+        private static bool IsNamedCommandWithArguments(Ast ast)
+        {
             CommandAst commandAst = ast as CommandAst;
 
                 return
                     commandAst != null &&
                     commandAst.InvocationOperator != TokenKind.Dot &&
                     PesterSymbolReference.GetCommandType(commandAst.GetCommandName()).HasValue &&
                     commandAst.CommandElements.Count >= 2;
-            },
-            true);
+        }
 
-            return commandAsts.Select(
-                ast =>
-                {
-                    // By this point we know the Ast is a CommandAst with 2 or more CommandElements
-                    int testNameParamIndex = 1;
-                    CommandAst testAst = (CommandAst)ast;
+        /// <summary>
+        /// Test whether the given CommandAst represents a Pester command
+        /// </summary>
+        /// <param name="commandAst">the CommandAst to test</param>
+        /// <returns>true if the CommandAst represents a Pester command, false otherwise</returns>
+        private static bool IsPesterCommand(CommandAst commandAst)
+        {
+            if (commandAst == null)
+            {
+                return false;
+            }
 
-                    // 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;
-                        }
-                    }
+            // Ensure the first word is a Pester keyword
+            if (!(commandAst.CommandElements[0] is StringConstantExpressionAst pesterKeywordAst &&
+                PesterSymbolReference.PesterKeywords.ContainsKey(pesterKeywordAst.Value)))
+            {
+                return false;
+            }
 
-                    if (testNameParamIndex > testAst.CommandElements.Count - 1)
-                    {
-                        return null;
-                    }
+            // Ensure that the last argument of the command is a scriptblock
+            if (!(commandAst.CommandElements[commandAst.CommandElements.Count-1] is ScriptBlockExpressionAst))
+            {
+                return false;
+            }
 
-                    StringConstantExpressionAst stringAst =
-                        testAst.CommandElements[testNameParamIndex] as StringConstantExpressionAst;
+            return true;
+        }
+
+        /// <summary>
+        /// Convert a CommandAst known to represent a Pester command and a reference to the scriptfile
+        /// it is in into symbol representing a Pester call for code lens
+        /// </summary>
+        /// <param name="scriptFile">the scriptfile the Pester call occurs in</param>
+        /// <param name="pesterCommandAst">the CommandAst representing the Pester call</param>
+        /// <returns>a symbol representing the Pester call containing metadata for CodeLens to use</returns>
+        private static PesterSymbolReference ConvertPesterAstToSymbolReference(ScriptFile scriptFile, CommandAst pesterCommandAst)
+        {
+            string testLine = scriptFile.GetLine(pesterCommandAst.Extent.StartLineNumber);
+            string commandName = (pesterCommandAst.CommandElements[0] as StringConstantExpressionAst)?.Value;
 
-                    if (stringAst == null)
+            // Search for a name for the test
+            string testName = null;
+            for (int i = 1; i < pesterCommandAst.CommandElements.Count; i++)
+            {
+                CommandElementAst currentCommandElement = pesterCommandAst.CommandElements[i];
+
+                // Check for an explicit "-Name" parameter
+                if (currentCommandElement is CommandParameterAst parameterAst)
+                {
+                    i++;
+                    if (parameterAst.ParameterName == "Name" && i < pesterCommandAst.CommandElements.Count)
                     {
-                        return null;
+                        testName = (pesterCommandAst.CommandElements[i] as StringConstantExpressionAst)?.Value;
+                        break;
                     }
+                    continue;
+                }
 
-                    string testDefinitionLine =
-                        scriptFile.GetLine(
-                            ast.Extent.StartLineNumber);
-
-                    return
-                        new PesterSymbolReference(
-                            scriptFile,
-                            testAst.GetCommandName(),
-                            testDefinitionLine,
-                            stringAst.Value,
-                            ast.Extent);
+                // Otherwise, if an argument is given with no parameter, we assume it's the name
+                if (pesterCommandAst.CommandElements[i] is StringConstantExpressionAst testNameStrAst)
+                {
+                    testName = testNameStrAst.Value;
+                    break;
+                }
+            }
 
-                }).Where(s => s != null);
+            return new PesterSymbolReference(
+                scriptFile,
+                commandName,
+                testLine,
+                testName,
+                pesterCommandAst.Extent
+            );
         }
     }
 
@@ -114,6 +155,17 @@ public enum PesterCommandType
     /// </summary>
     public class PesterSymbolReference : SymbolReference
     {
+        /// <summary>
+        /// Lookup for Pester keywords we support. Ideally we could extract these from Pester itself
+        /// </summary>
+        internal static readonly IReadOnlyDictionary<string, PesterCommandType> PesterKeywords =
+            new Dictionary<string, PesterCommandType>(StringComparer.OrdinalIgnoreCase)
+            {
+                { "Describe", PesterCommandType.Describe },
+                { "Context", PesterCommandType.Context },
+                { "It", PesterCommandType.It }
+            };
+
         private static char[] DefinitionTrimChars = new char[] { ' ', '{' };
 
         /// <summary>
@@ -145,25 +197,12 @@ internal PesterSymbolReference(
 
         internal static PesterCommandType? GetCommandType(string commandName)
         {
-            if (commandName == null)
+            PesterCommandType pesterCommandType;
+            if (!PesterKeywords.TryGetValue(commandName, out pesterCommandType))
             {
                 return null;
             }
-
-            switch (commandName.ToLower())
-            {
-                case "describe":
-                    return PesterCommandType.Describe;
-
-                case "context":
-                    return PesterCommandType.Context;
-
-                case "it":
-                    return PesterCommandType.It;
-
-                default:
-                    return null;
-            }
+            return pesterCommandType;
         }
     }
 }