Create a tool to provide prediction to Az commands

tools/Az.Tools.AzPredictor/AzPredictor.Test/AzPredictor.Test.csproj

@@ -0,0 +1,34 @@
+<Project Sdk="Microsoft.NET.Sdk">
+
+  <PropertyGroup>
+    <TargetFramework>net5.0</TargetFramework>
+    <AssemblyName>Microsoft.Azure.PowerShell.AzPredictor.Test</AssemblyName>
+    <RootNamespace>Microsoft.Azure.PowerShell.AzPredictor.Test</RootNamespace>
+    <IsPackable>false</IsPackable>
+    <GenerateAssemblyInfo>false</GenerateAssemblyInfo>
+    <TreatWarningsAsErrors>true</TreatWarningsAsErrors>
+    <DocumentationFile>$(OutputPath)\Microsoft.Azure.PowerShell.AzPredictor.Test.xml</DocumentationFile>
+    <GenerateSerializationAssemblies>Auto</GenerateSerializationAssemblies>
+  </PropertyGroup>
+
+  <ItemGroup>
+    <PackageReference Include="System.Management.Automation" Version="7.1.0-preview.9" />
+    <PackageReference Include="Microsoft.NET.Test.Sdk" Version="16.6.1" />
+    <PackageReference Include="xunit" Version="2.4.1" />
+    <PackageReference Include="xunit.runner.visualstudio" Version="2.4.3">
+      <IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
+      <PrivateAssets>all</PrivateAssets>
+    </PackageReference>
+    <PackageReference Include="coverlet.collector" Version="1.3.0">
+      <IncludeAssets>runtime; build; native; contentfiles; analyzers; buildtransitive</IncludeAssets>
+      <PrivateAssets>all</PrivateAssets>
+    </PackageReference>
+    <ProjectReference Include="..\AzPredictor\AzPredictor.csproj" />
+  </ItemGroup>
+
+  <ItemGroup>
+    <None Include="Data\CommandsModel.zip" CopyToOutputDirectory="PreserveNewest" />
+    <None Include="Data\PredictionsModel.zip" CopyToOutputDirectory="PreserveNewest" />
+  </ItemGroup>
+
+</Project>

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

@@ -0,0 +1,104 @@
+// ----------------------------------------------------------------------------------
+//
+// Copyright Microsoft Corporation
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ----------------------------------------------------------------------------------
+
+using Microsoft.Azure.PowerShell.AzPredictor.Test.Mocks;
+using System.Management.Automation.Subsystem;
+using System.Threading;
+using Xunit;
+
+namespace Microsoft.Azure.PowerShell.AzPredictor.Test
+{
+    /// <summary>
+    /// Tests for <see cref="AzPredictorService"/>
+    /// </summary>
+    [Collection("Model collection")]
+    public class AzPredictorServiceTests
+    {
+        private readonly ModelFixture _fixture;
+        private readonly AzPredictorService _service;
+        private readonly Predictor _suggestionsPredictor;
+        private readonly Predictor _commandsPredictor;
+
+        /// <summary>
+        /// Constructs a new instance of <see cref="AzPredictorServiceTests"/>
+        /// </summary>
+        /// <param name="fixture"></param>
+        public AzPredictorServiceTests(ModelFixture fixture)
+        {
+            this._fixture = fixture;
+            var startHistory = $"{AzPredictorConstants.CommandHistoryPlaceholder}{AzPredictorConstants.CommandConcatenator}{AzPredictorConstants.CommandHistoryPlaceholder}";
+            this._suggestionsPredictor = new Predictor(this._fixture.PredictionCollection[startHistory]);
+            this._commandsPredictor = new Predictor(this._fixture.CommandCollection);
+
+            this._service = new MockAzPredictorService(this._fixture.PredictionCollection[startHistory], this._fixture.CommandCollection);
+        }
+
+
+        /// <summary>
+        /// Verifies that the prediction comes from the suggestions list, not the command list.
+        /// </summary>
+        [Theory]
+        [InlineData("CONNECT-AZACCOUNT")]
+        [InlineData("set-azstorageaccount ")]
+        [InlineData("Get-AzResourceG")]
+        [InlineData("Get-AzStorageAcco")]
+        [InlineData("Get-AzKeyVault -VaultName")]
+        [InlineData("GET-AZSTORAGEACCOUNTKEY -NAME ")]
+        [InlineData("new-azresourcegroup -name hello")]
+        [InlineData("Get-AzContext -Name")]
+        [InlineData("Get-AzContext -ErrorAction")]
+        public void VerifyUsingSuggestion(string userInput)
+        {
+            var predictionContext = PredictionContext.Create(userInput);
+            var expected = this._suggestionsPredictor.Query(predictionContext.InputAst, CancellationToken.None);
+            var actual = this._service.GetSuggestion(predictionContext.InputAst, CancellationToken.None);
+            Assert.Equal(expected, actual);
+            Assert.NotNull(actual);
+        }
+
+
+        /// <summary>
+        /// Verifies that when no prediction is in the suggestion list, we'll use the command list.
+        /// </summary>
+        [Theory]
+        [InlineData("Get-AzResource -Name hello -Pre")]
+        [InlineData("Get-AzADServicePrincipal -ApplicationObject")]
+        public void VerifyUsingCommand(string userInput)
+        {
+            var predictionContext = PredictionContext.Create(userInput);
+            var expected = this._commandsPredictor.Query(predictionContext.InputAst, CancellationToken.None);
+            var actual = this._service.GetSuggestion(predictionContext.InputAst, CancellationToken.None);
+            Assert.Equal(expected, actual);
+            Assert.NotNull(actual);
+        }
+
+        /// <summary>
+        /// Verify that no prediction for the user input, meaning it's not in the prediction list or the command list.
+        /// </summary>
+        [Theory]
+        [InlineData(AzPredictorConstants.CommandHistoryPlaceholder)]
+        [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, CancellationToken.None);
+            Assert.Null(actual);
+        }
+    }
+}

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

@@ -0,0 +1,155 @@
+// ----------------------------------------------------------------------------------
+//
+// Copyright Microsoft Corporation
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ----------------------------------------------------------------------------------
+
+using Microsoft.Azure.PowerShell.AzPredictor.Test.Mocks;
+using System.Collections.Generic;
+using System.Linq;
+using System.Management.Automation.Subsystem;
+using System.Threading;
+using Xunit;
+
+namespace Microsoft.Azure.PowerShell.AzPredictor.Test
+{
+    /// <summary>
+    /// Tests for <see cref="AzPredictor"/>
+    /// </summary>
+    [Collection("Model collection")]
+    public sealed class AzPredictorTests
+    {
+        private readonly ModelFixture _fixture;
+        private readonly MockAzPredictorTelemetryClient _telemetryClient;
+        private readonly MockAzPredictorService _service;
+        private readonly AzPredictor _azPredictor;
+
+        /// <summary>
+        /// Constructs a new instance of <see cref="AzPredictorTests"/>
+        /// </summary>
+        public AzPredictorTests(ModelFixture modelFixture)
+        {
+            this._fixture = modelFixture;
+            var startHistory = $"{AzPredictorConstants.CommandHistoryPlaceholder}{AzPredictorConstants.CommandConcatenator}{AzPredictorConstants.CommandHistoryPlaceholder}";
+
+            this._service = new MockAzPredictorService(this._fixture.PredictionCollection[startHistory], this._fixture.CommandCollection);
+            this._telemetryClient = new MockAzPredictorTelemetryClient();
+            this._azPredictor = new AzPredictor(this._service, this._telemetryClient);
+        }
+
+        /// <summary>
+        /// Verifies when the last command in history are not supported.
+        /// We don't collect the telemetry and only request prediction while StartEarlyProcess is called.
+        /// </summary>
+        [Theory]
+        [InlineData("start_of_snippet\nstart_of_snippet\nstart_of_snippet")]
+        [InlineData("start_of_snippet")]
+        [InlineData("")]
+        [InlineData("git status")]
+        [InlineData("git status\nGet-ChildItem")]
+        [InlineData("^29a9l2")]
+        [InlineData("'Get-AzResource'")]
+        [InlineData("Get-AzResource\ngit log")]
+        [InlineData("Get-ChildItem")]
+        public void VerifyWithNonSupportedCommand(string historyLine)
+        {
+            IReadOnlyList<string> history = historyLine.Split('\n');
+
+            this._telemetryClient.RecordedSuggestion = null;
+            this._service.IsPredictionRequested = false;
+
+            this._azPredictor.StartEarlyProcessing(history);
+
+            Assert.True(this._service.IsPredictionRequested);
+            Assert.Null(this._telemetryClient.RecordedSuggestion);
+        }
+
+        /// <summary>
+        /// Verifies when the last command in history are not supported.
+        /// We don't collect the telemetry and only request prediction while StartEarlyProcess is called.
+        /// </summary>
+        [Theory]
+        [InlineData("start_of_snippet\nConnect-AzAccount")]
+        [InlineData("Get-AzResource")]
+        [InlineData("git status\nGet-AzContext")]
+        [InlineData("Get-AzContext\nGet-AzLog")]
+        public void VerifyWithOneSupportedCommand(string historyLine)
+        {
+            IReadOnlyList<string> history = historyLine.Split('\n');
+
+            this._telemetryClient.RecordedSuggestion = null;
+            this._service.IsPredictionRequested = false;
+
+            this._azPredictor.StartEarlyProcessing(history);
+
+            Assert.True(this._service.IsPredictionRequested);
+            Assert.NotNull(this._telemetryClient.RecordedSuggestion);
+        }
+
+        /// <summary>
+        /// Verify that the supported commands parameter values are masked.
+        /// </summary>
+        [Fact]
+        public void VerifySupportedCommandMasked()
+        {
+            var input = "Get-AzVMExtension -ResourceGroupName 'ResourceGroup11' -VMName 'VirtualMachine22'";
+            var expected = "Get-AzVMExtension -ResourceGroupName *** -VMName ***";
+
+            this._telemetryClient.RecordedSuggestion = null;
+            this._service.IsPredictionRequested = false;
+
+            this._azPredictor.StartEarlyProcessing(new List<string> { input } );
+
+            Assert.True(this._service.IsPredictionRequested);
+            Assert.NotNull(this._telemetryClient.RecordedSuggestion);
+            Assert.Equal(expected, this._telemetryClient.RecordedSuggestion.HistoryLine);
+
+            input = "Get-AzStorageAccountKey -Name:'ContosoStorage' -ResourceGroupName:'ContosoGroup02'";
+            expected = "Get-AzStorageAccountKey -Name:*** -ResourceGroupName:***";
+
+
+            this._telemetryClient.RecordedSuggestion = null;
+            this._service.IsPredictionRequested = false;
+
+            this._azPredictor.StartEarlyProcessing(new List<string> { input } );
+
+            Assert.True(this._service.IsPredictionRequested);
+            Assert.NotNull(this._telemetryClient.RecordedSuggestion);
+            Assert.Equal(expected, this._telemetryClient.RecordedSuggestion.HistoryLine);
+        }
+
+        /// <summary>
+        /// 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")]
+        [InlineData("Get-AzADServicePrincipal -ApplicationObject")]
+        public void VerifySuggestion(string userInput)
+        {
+            var predictionContext = PredictionContext.Create(userInput);
+            var expected = this._service.GetSuggestion(predictionContext.InputAst, CancellationToken.None);
+            var actual = this._azPredictor.GetSuggestion(predictionContext, CancellationToken.None);
+            if (actual == null)
+            {
+                Assert.Null(expected);
+            }
+            else
+            {
+                Assert.Single(actual);
+                Assert.Equal(expected, actual.First().SuggestionText);
+            }
+
+        }
+    }
+}

tools/Az.Tools.AzPredictor/AzPredictor.Test/Mocks/MockAzPredictorService.cs

@@ -0,0 +1,52 @@
+// ----------------------------------------------------------------------------------
+//
+// Copyright Microsoft Corporation
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ----------------------------------------------------------------------------------
+
+using System.Collections.Generic;
+
+namespace Microsoft.Azure.PowerShell.AzPredictor.Test.Mocks
+{
+    /// <summary>
+    /// Mock <see cref="AzPredictorService"/> so that it doesn't do httpd request to get the commands and predictions.
+    /// </summary>
+    public sealed class MockAzPredictorService : AzPredictorService
+    {
+        /// <summary>
+        /// Gets or sets if a predictions is requested.
+        /// </summary>
+        public bool IsPredictionRequested { get; set; }
+
+        /// <summary>
+        /// Constructs a new instance of <see cref="MockAzPredictorService"/>
+        /// </summary>
+        /// <param name="suggestions">The suggestions collection</param>
+        /// <param name="commands">The commands collection</param>
+        public MockAzPredictorService(IList<string> suggestions, IList<string> commands)
+        {
+            SetCommandsPredictor(commands);
+            SetSuggestionPredictor(suggestions);
+        }
+
+        /// <inheritdoc/>
+        public override void RequestPredictions(string history)
+        {
+            this.IsPredictionRequested = true;
+        }
+
+        /// <inheritdoc/>
+        protected override void RequestCommands()
+        {
+            // Do nothing since we've set the command and suggestion predictors.
+        }
+    }
+}

tools/Az.Tools.AzPredictor/AzPredictor.Test/Mocks/MockAzPredictorTelemetryClient.cs

@@ -0,0 +1,50 @@
+// ----------------------------------------------------------------------------------
+//
+// Copyright Microsoft Corporation
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ----------------------------------------------------------------------------------
+
+using System.Collections.Generic;
+
+namespace Microsoft.Azure.PowerShell.AzPredictor.Test.Mocks
+{
+    sealed class MockAzPredictorTelemetryClient : ITelemetryClient
+    {
+        public class RecordedSuggestionForHistory
+        {
+            public string HistoryLine { get; set; }
+            public int? SuggestionIndex { get; set; }
+            public int? FallbackIndex { get; set; }
+            public IEnumerable<string> TopSuggestions { get; set; }
+        }
+
+        public RecordedSuggestionForHistory RecordedSuggestion { get; set; }
+        public int SuggestionAccepted { get; set; }
+
+        /// <inheritdoc/>
+        public void OnSuggestionForHistory(string historyLine, int? suggestionIndex, int? fallbackIndex, IEnumerable<string> topSuggestions)
+        {
+            this.RecordedSuggestion = new RecordedSuggestionForHistory()
+            {
+                HistoryLine = historyLine,
+                SuggestionIndex = suggestionIndex,
+                FallbackIndex = fallbackIndex,
+                TopSuggestions = topSuggestions
+            };
+        }
+
+        /// <inheritdoc/>
+        public void OnSuggestionAccepted(string acceptedSuggestion)
+        {
+            ++this.SuggestionAccepted;
+        }
+    }
+}