Az.Tools.Predictor refactor and performance improvement.

tools/Az.Tools.Predictor/Az.Tools.Predictor.Test/AzPredictorServiceTests.cs

@@ -1,4 +1,4 @@
-// ----------------------------------------------------------------------------------
+// ----------------------------------------------------------------------------------
 //
 // Copyright Microsoft Corporation
 // Licensed under the Apache License, Version 2.0 (the "License");
@@ -13,7 +13,10 @@
 // ----------------------------------------------------------------------------------
 
 using Microsoft.Azure.PowerShell.Tools.AzPredictor.Test.Mocks;
+using System;
+using System.Collections.Generic;
 using System.Linq;
+using System.Management.Automation.Language;
 using System.Management.Automation.Subsystem;
 using System.Threading;
 using Xunit;
@@ -26,10 +29,36 @@ namespace Microsoft.Azure.PowerShell.Tools.AzPredictor.Test
     [Collection("Model collection")]
     public class AzPredictorServiceTests
     {
+        private class PredictiveSuggestionComparer : EqualityComparer<PredictiveSuggestion>
+        {
+            public override bool Equals(PredictiveSuggestion first, PredictiveSuggestion second)
+            {
+                if ((first == null) && (second == null))
+                {
+                    return true;
+                }
+                else if ((first == null) || (second == null))
+                {
+                    return false;
+                }
+
+                return string.Equals(first.SuggestionText, second.SuggestionText, StringComparison.Ordinal);
+            }
+
+            public override int GetHashCode(PredictiveSuggestion suggestion)
+            {
+                return suggestion.SuggestionText.GetHashCode();
+            }
+        }
+
         private readonly ModelFixture _fixture;
         private readonly AzPredictorService _service;
-        private readonly Predictor _suggestionsPredictor;
-        private readonly Predictor _commandsPredictor;
+        private readonly CommandLinePredictor _commandBasedPredictor;
+        private readonly CommandLinePredictor _fallbackPredictor;
+
+        private readonly AzPredictorService _noFallbackPredictorService;
+        private readonly AzPredictorService _noCommandBasedPredictorService;
+        private readonly AzPredictorService _noPredictorService;
 
         /// <summary>
         /// Constructs a new instance of <see cref="AzPredictorServiceTests"/>
@@ -39,15 +68,36 @@ public AzPredictorServiceTests(ModelFixture fixture)
         {
             this._fixture = fixture;
             var startHistory = $"{AzPredictorConstants.CommandPlaceholder}{AzPredictorConstants.CommandConcatenator}{AzPredictorConstants.CommandPlaceholder}";
-            this._suggestionsPredictor = new Predictor(this._fixture.PredictionCollection[startHistory], null);
-            this._commandsPredictor = new Predictor(this._fixture.CommandCollection, null);
+            this._commandBasedPredictor = new CommandLinePredictor(this._fixture.PredictionCollection[startHistory], null);
+            this._fallbackPredictor = new CommandLinePredictor(this._fixture.CommandCollection, null);
 
             this._service = new MockAzPredictorService(startHistory, this._fixture.PredictionCollection[startHistory], this._fixture.CommandCollection);
+
+            this._noFallbackPredictorService = new MockAzPredictorService(startHistory, this._fixture.PredictionCollection[startHistory], null);
+            this._noCommandBasedPredictorService = new MockAzPredictorService(null, null, this._fixture.CommandCollection);
+            this._noPredictorService = new MockAzPredictorService(null, null, null);
         }
 
+        /// <summary>
+        /// Verify the method checks parameter values.
+        /// </summary>
+        [Fact]
+        public void VerifyParameterValues()
+        {
+            var predictionContext = PredictionContext.Create("Get-AzContext");
+
+            Action actual = () => this._service.GetSuggestion(null, 1, 1, CancellationToken.None);
+            Assert.Throws<ArgumentNullException>(actual);
+
+            actual = () => this._service.GetSuggestion(predictionContext.InputAst, 0, 1, CancellationToken.None);
+            Assert.Throws<ArgumentOutOfRangeException>(actual);
+
+            actual = () => this._service.GetSuggestion(predictionContext.InputAst, 1, 0, CancellationToken.None);
+            Assert.Throws<ArgumentOutOfRangeException>(actual);
+        }
 
         /// <summary>
-        /// Verifies that the prediction comes from the suggestions list, not the command list.
+        /// Verifies that the prediction comes from the command based list, not the fallback list.
         /// </summary>
         [Theory]
         [InlineData("CONNECT-AZACCOUNT")]
@@ -59,52 +109,126 @@ public AzPredictorServiceTests(ModelFixture fixture)
         [InlineData("new-azresourcegroup -name hello")]
         [InlineData("Get-AzContext -Name")]
         [InlineData("Get-AzContext -ErrorAction")]
-        public void VerifyUsingSuggestion(string userInput)
+        public void VerifyUsingCommandBasedPredictor(string userInput)
         {
             var predictionContext = PredictionContext.Create(userInput);
-            var presentCommands = new System.Collections.Generic.Dictionary<string, int>();
-            var expected = this._suggestionsPredictor.Query(predictionContext.InputAst, presentCommands, 1, 1, CancellationToken.None);
+            var commandAst = predictionContext.InputAst.FindAll(p => p is CommandAst, true).LastOrDefault() as CommandAst;
+            var commandName = (commandAst?.CommandElements?.FirstOrDefault() as StringConstantExpressionAst)?.Value;
+            var inputParameterSet = new ParameterSet(commandAst);
+            var rawUserInput = predictionContext.InputAst.Extent.Text;
+            var presentCommands = new Dictionary<string, int>();
+            var expected = this._commandBasedPredictor.GetSuggestion(commandName,
+                    inputParameterSet,
+                    rawUserInput,
+                    presentCommands,
+                    1,
+                    1,
+                    CancellationToken.None);
+
             var actual = this._service.GetSuggestion(predictionContext.InputAst, 1, 1, CancellationToken.None);
-            Assert.NotEmpty(actual);
-            Assert.NotNull(actual.First().Item1);
-            Assert.Equal(expected.Item1.First().Key, actual.First().Item1);
-            Assert.Equal(PredictionSource.CurrentCommand, actual.First().Item3);
+            Assert.NotNull(actual);
+            Assert.True(actual.Count > 0);
+            Assert.NotNull(actual.PredictiveSuggestions.First());
+            Assert.NotNull(actual.PredictiveSuggestions.First().SuggestionText);
+            Assert.Equal(expected.Count, actual.Count);
+            Assert.Equal<PredictiveSuggestion>(expected.PredictiveSuggestions, actual.PredictiveSuggestions, new PredictiveSuggestionComparer());
+            Assert.Equal<string>(expected.SourceTexts, actual.SourceTexts);
+            Assert.All<SuggestionSource>(actual.SuggestionSources, (source) => Assert.Equal(SuggestionSource.CurrentCommand, source));
+
+            actual = this._noFallbackPredictorService.GetSuggestion(predictionContext.InputAst, 1, 1, CancellationToken.None);
+            Assert.NotNull(actual);
+            Assert.True(actual.Count > 0);
+            Assert.NotNull(actual.PredictiveSuggestions.First());
+            Assert.NotNull(actual.PredictiveSuggestions.First().SuggestionText);
+            Assert.Equal(expected.Count, actual.Count);
+            Assert.Equal<PredictiveSuggestion>(expected.PredictiveSuggestions, actual.PredictiveSuggestions, new PredictiveSuggestionComparer());
+            Assert.Equal<string>(expected.SourceTexts, actual.SourceTexts);
+            Assert.All<SuggestionSource>(actual.SuggestionSources, (source) => Assert.Equal(SuggestionSource.CurrentCommand, source));
         }
 
         /// <summary>
-        /// Verifies that when no prediction is in the suggestion list, we'll use the command list.
+        /// Verifies that when no prediction is in the command based list, we'll use the fallback list.
         /// </summary>
         [Theory]
         [InlineData("Get-AzResource -Name hello -Pre")]
         [InlineData("Get-AzADServicePrincipal -ApplicationObject")]
-        public void VerifyUsingCommand(string userInput)
+        public void VerifyUsingFallbackPredictor(string userInput)
         {
             var predictionContext = PredictionContext.Create(userInput);
-            var presentCommands = new System.Collections.Generic.Dictionary<string, int>();
-            var expected = this._commandsPredictor.Query(predictionContext.InputAst, presentCommands, 1, 1, CancellationToken.None);
+            var commandAst = predictionContext.InputAst.FindAll(p => p is CommandAst, true).LastOrDefault() as CommandAst;
+            var commandName = (commandAst?.CommandElements?.FirstOrDefault() as StringConstantExpressionAst)?.Value;
+            var inputParameterSet = new ParameterSet(commandAst);
+            var rawUserInput = predictionContext.InputAst.Extent.Text;
+            var presentCommands = new Dictionary<string, int>();
+            var expected = this._fallbackPredictor.GetSuggestion(commandName,
+                    inputParameterSet,
+                    rawUserInput,
+                    presentCommands,
+                    1,
+                    1,
+                    CancellationToken.None);
+
             var actual = this._service.GetSuggestion(predictionContext.InputAst, 1, 1, CancellationToken.None);
-            Assert.NotEmpty(actual);
-            Assert.NotNull(actual.First().Item1);
-            Assert.Equal(expected.Item1.First().Key, actual.First().Item1);
-            Assert.Equal(PredictionSource.StaticCommands, actual.First().Item3);
+            Assert.NotNull(actual);
+            Assert.True(actual.Count > 0);
+            Assert.NotNull(actual.PredictiveSuggestions.First());
+            Assert.NotNull(actual.PredictiveSuggestions.First().SuggestionText);
+            Assert.Equal(expected.Count, actual.Count);
+            Assert.Equal<PredictiveSuggestion>(expected.PredictiveSuggestions, actual.PredictiveSuggestions, new PredictiveSuggestionComparer());
+            Assert.Equal<string>(expected.SourceTexts, actual.SourceTexts);
+            Assert.All<SuggestionSource>(actual.SuggestionSources, (source) => Assert.Equal(SuggestionSource.StaticCommands, source));
+
+            actual = this._noCommandBasedPredictorService.GetSuggestion(predictionContext.InputAst, 1, 1, CancellationToken.None);
+            Assert.NotNull(actual);
+            Assert.True(actual.Count > 0);
+            Assert.NotNull(actual.PredictiveSuggestions.First());
+            Assert.NotNull(actual.PredictiveSuggestions.First().SuggestionText);
+            Assert.Equal(expected.Count, actual.Count);
+            Assert.Equal<PredictiveSuggestion>(expected.PredictiveSuggestions, actual.PredictiveSuggestions, new PredictiveSuggestionComparer());
+            Assert.Equal<string>(expected.SourceTexts, actual.SourceTexts);
+            Assert.All<SuggestionSource>(actual.SuggestionSources, (source) => Assert.Equal(SuggestionSource.StaticCommands, source));
         }
 
         /// <summary>
-        /// Verify that no prediction for the user input, meaning it's not in the prediction list or the command list.
+        /// Verify that no prediction for the user input, meaning it's not in the command based list or the fallback list.
         /// </summary>
         [Theory]
         [InlineData(AzPredictorConstants.CommandPlaceholder)]
-        [InlineData("git status")]
         [InlineData("Get-ChildItem")]
         [InlineData("new-azresourcegroup -NoExistingParam")]
         [InlineData("get-azaccount ")]
-        [InlineData("Get-AzContext Name")]
         [InlineData("NEW-AZCONTEXT")]
         public void VerifyNoPrediction(string userInput)
         {
             var predictionContext = PredictionContext.Create(userInput);
             var actual = this._service.GetSuggestion(predictionContext.InputAst, 1, 1, CancellationToken.None);
-            Assert.Empty(actual);
+            Assert.Equal(0, actual.Count);
+
+            actual = this._noFallbackPredictorService.GetSuggestion(predictionContext.InputAst, 1, 1, CancellationToken.None);
+            Assert.Equal(0, actual.Count);
+
+            actual = this._noCommandBasedPredictorService.GetSuggestion(predictionContext.InputAst, 1, 1, CancellationToken.None);
+            Assert.Equal(0, actual.Count);
+
+            actual = this._noPredictorService.GetSuggestion(predictionContext.InputAst, 1, 1, CancellationToken.None);
+            Assert.Null(actual);
+        }
+
+        /// <summary>
+        /// Verify when we cannot parse the user input correctly.
+        /// </summary>
+        /// <remarks>
+        /// When we can parse them correctly, please move the InlineData to the corresponding test methods, for example, "git status"
+        /// doesn't have any prediction so it should move to <see cref="VerifyNoPrediction"/>.
+        /// </remarks>
+        [Theory]
+        [InlineData("git status")]
+        [InlineData("Get-AzContext Name")]
+        public void VerifyMalFormattedCommandLine(string userInput)
+        {
+            var predictionContext = PredictionContext.Create(userInput);
+            Action actual = () => this._service.GetSuggestion(predictionContext.InputAst, 1, 1, CancellationToken.None);
+            _ = Assert.Throws<InvalidOperationException>(actual);
         }
     }
 }

tools/Az.Tools.Predictor/Az.Tools.Predictor.Test/AzPredictorTests.cs

@@ -45,6 +45,7 @@ public AzPredictorTests(ModelFixture modelFixture)
             this._azPredictor = new AzPredictor(this._service, this._telemetryClient, new Settings()
             {
                 SuggestionCount = 1,
+                MaxAllowedCommandDuplicate = 1,
             },
             null);
         }
@@ -134,7 +135,6 @@ public void VerifySupportedCommandMasked()
         /// Verifies AzPredictor returns the same value as AzPredictorService for the prediction.
         /// </summary>
         [Theory]
-        [InlineData("git status")]
         [InlineData("new-azresourcegroup -name hello")]
         [InlineData("Get-AzContext -Name")]
         [InlineData("Get-AzContext -ErrorAction")]
@@ -145,7 +145,8 @@ public void VerifySuggestion(string userInput)
             var expected = this._service.GetSuggestion(predictionContext.InputAst, 1, 1, CancellationToken.None);
             var actual = this._azPredictor.GetSuggestion(predictionContext, CancellationToken.None);
 
-            Assert.Equal(expected.Select(e => e.Item1), actual.Select(a => a.SuggestionText));
+            Assert.Equal(expected.Count, actual.Count);
+            Assert.Equal(expected.PredictiveSuggestions.First().SuggestionText, actual.First().SuggestionText);
         }
 
         /// <summary>
@@ -158,6 +159,7 @@ public void VerifySuggestionOnIncompleteCommand()
             var localAzPredictor = new AzPredictor(this._service, this._telemetryClient, new Settings()
             {
                 SuggestionCount = 7,
+                MaxAllowedCommandDuplicate = 1,
             },
             null);
 
@@ -169,5 +171,23 @@ public void VerifySuggestionOnIncompleteCommand()
 
             Assert.Equal(expected, actual.First().SuggestionText);
         }
+
+
+        /// <summary>
+        /// Verify when we cannot parse the user input correctly.
+        /// </summary>
+        /// <remarks>
+        /// When we can parse them correctly, please move the InlineData to the corresponding test methods, for example, "git status"
+        /// can be moved to <see cref="VerifySuggestion"/>.
+        /// </remarks>
+        [Theory]
+        [InlineData("git status")]
+        public void VerifyMalFormattedCommandLine(string userInput)
+        {
+            var predictionContext = PredictionContext.Create(userInput);
+            var actual = this._azPredictor.GetSuggestion(predictionContext, CancellationToken.None);
+
+            Assert.Empty(actual);
+        }
     }
 }