Fix the unit tests

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

@@ -62,7 +62,7 @@ public override int GetHashCode(PredictiveSuggestion suggestion)
         private readonly AzPredictorService _noPredictorService;
         private readonly AzContext _azContext;
 
-        private PowerShellRuntime _powerShellRuntime;
+        private MockPowerShellRuntime _powerShellRuntime;
 
         /// <summary>
         /// Constructs a new instance of <see cref="AzPredictorServiceTests"/>
@@ -71,7 +71,7 @@ public override int GetHashCode(PredictiveSuggestion suggestion)
         public AzPredictorServiceTests(ModelFixture fixture)
         {
             this._fixture = fixture;
-            _powerShellRuntime = new PowerShellRuntime();
+            _powerShellRuntime = new MockPowerShellRuntime();
             _azContext = new AzContext(_powerShellRuntime);
             var startHistory = $"{AzPredictorConstants.CommandPlaceholder}{AzPredictorConstants.CommandConcatenator}{AzPredictorConstants.CommandPlaceholder}";
             this._commandBasedPredictor = new CommandLinePredictor(this._fixture.PredictionCollection[startHistory], null, null, _azContext);

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

@@ -12,6 +12,7 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
+using Microsoft.Azure.PowerShell.Tools.AzPredictor.Test.Mocks;
 using Microsoft.Azure.PowerShell.Tools.AzPredictor.Utilities;
 using System;
 using System.Collections.Generic;
@@ -32,15 +33,15 @@ public sealed class CommandLinePredictorTests : IDisposable
         private readonly ModelFixture _fixture;
         private readonly CommandLinePredictor _predictor;
         private readonly AzContext _azContext;
-        private PowerShellRuntime _powerShellRuntime;
+        private MockPowerShellRuntime _powerShellRuntime;
 
         /// <summary>
         /// Constructs a new instance of <see cref="CommandLinePredictorTests" />
         /// </summary>
         public CommandLinePredictorTests(ModelFixture fixture)
         {
             _fixture = fixture;
-            _powerShellRuntime = new PowerShellRuntime();
+            _powerShellRuntime = new MockPowerShellRuntime();
             _azContext = new AzContext(_powerShellRuntime);
             var startHistory = $"{AzPredictorConstants.CommandPlaceholder}{AzPredictorConstants.CommandConcatenator}{AzPredictorConstants.CommandPlaceholder}";
             _predictor = new CommandLinePredictor(_fixture.PredictionCollection[startHistory], null,null,  _azContext);

tools/Az.Tools.Predictor/Az.Tools.Predictor.Test/Mocks/MockPowerShellRuntime.cs

@@ -0,0 +1,47 @@
+// ----------------------------------------------------------------------------------
+//
+// 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.Tools.AzPredictor.Utilities;
+using System;
+using System.Collections.Generic;
+using System.Management.Automation.Runspaces;
+
+namespace Microsoft.Azure.PowerShell.Tools.AzPredictor.Test.Mocks
+{
+    using PowerShell = System.Management.Automation.PowerShell;
+
+    /// <summary>
+    /// A Mock PowerShell environment to be used in test cases.
+    /// </summary>
+    internal sealed class MockPowerShellRuntime : IPowerShellRuntime, IDisposable
+    {
+        /// <inheritdoc />
+        public Runspace DefaultRunspace { get; private set; } = PowerShellRunspaceUtilities.GetMinimalRunspace();
+
+        /// <inheritdoc />
+        public PowerShell ConsoleRuntime => throw new NotImplementedException("It's not implemented yet because there is no test case to set up powershell environment.");
+
+        /// <inheritdoc />
+        public IList<T> ExecuteScript<T>(string contents) => throw new NotImplementedException("It's not implemented yet because there is no test case to set up powershell environment.");
+
+        public void Dispose()
+        {
+            if (DefaultRunspace is not null)
+            {
+                DefaultRunspace.Dispose();
+                DefaultRunspace = null;
+            }
+        }
+    }
+}

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

@@ -12,6 +12,7 @@
 // limitations under the License.
 // ----------------------------------------------------------------------------------
 
+using Microsoft.Azure.PowerShell.Tools.AzPredictor.Test.Mocks;
 using Microsoft.Azure.PowerShell.Tools.AzPredictor.Utilities;
 using System;
 using System.Collections.Generic;
@@ -28,14 +29,14 @@ namespace Microsoft.Azure.PowerShell.Tools.AzPredictor.Test
     public sealed class ParameterSetTests : IDisposable
     {
         private readonly AzContext _azContext;
-        private PowerShellRuntime _powerShellRuntime;
+        private MockPowerShellRuntime _powerShellRuntime;
 
         /// <summary>
         /// Creates a new instance of <see cref="ParameterSetTests" />.
         /// </summary>
         public ParameterSetTests()
         {
-            _powerShellRuntime = new PowerShellRuntime();
+            _powerShellRuntime = new MockPowerShellRuntime();
             _azContext = new AzContext(_powerShellRuntime);
         }
 

tools/Az.Tools.Predictor/Az.Tools.Predictor/AzContext.cs

@@ -32,7 +32,7 @@ internal sealed class AzContext : IAzContext
     {
         private const string InternalUserSuffix = "@microsoft.com";
         private static readonly Version DefaultVersion = new Version("0.0.0.0");
-        private PowerShellRuntime _powerShellRuntime;
+        private IPowerShellRuntime _powerShellRuntime;
 
         /// <inheritdoc/>
         public Version AzVersion { get; private set; } = DefaultVersion;
@@ -129,7 +129,7 @@ public Version ModuleVersion
         /// <inheritdoc/>
         public bool IsInternal { get; internal set; }
 
-        public AzContext(PowerShellRuntime powerShellRuntime) => _powerShellRuntime
+        public AzContext(IPowerShellRuntime powerShellRuntime) => _powerShellRuntime
              = powerShellRuntime;
 
         /// <inheritdoc/>

tools/Az.Tools.Predictor/Az.Tools.Predictor/AzPredictorSurveyHelper.cs

@@ -35,7 +35,7 @@ internal sealed class AzPredictorSurveyHelper : ISurveyHelper, IDisposable
         private DateTime _lastCheckedTime = DateTime.MinValue;
         private Timer _promptDelayTimer;
 
-        public AzPredictorSurveyHelper(PowerShellRuntime powerShellRuntime)
+        public AzPredictorSurveyHelper(IPowerShellRuntime powerShellRuntime)
         {
             var promptMessageScript = @"
                 if ([Microsoft.Azure.PowerShell.Tools.AzPredictor.AzPredictorData]::ShowSurveyOnIdle) {

tools/Az.Tools.Predictor/Az.Tools.Predictor/Utilities/IPowerShellRuntime.cs

@@ -0,0 +1,48 @@
+// ----------------------------------------------------------------------------------
+//
+// 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;
+using System.Management.Automation.Runspaces;
+
+namespace Microsoft.Azure.PowerShell.Tools.AzPredictor.Utilities
+{
+    using PowerShell = System.Management.Automation.PowerShell;
+
+    /// <summary>
+    /// A PowerShell environment to run PowerShell cmdlets and scripts.
+    /// </summary>
+    internal interface IPowerShellRuntime
+    {
+        /// <summary>
+        /// Gets the minimum PowerShell Runspace. This isn't the necessary the same one as the PowerShell environment that Az
+        /// Predictor is running on.
+        /// </summary>
+        Runspace DefaultRunspace { get; }
+
+        /// <summary>
+        /// The PowerShell environment that the module is imported into.
+        /// </summary>
+        /// <remarks>
+        /// The usage of <see cref="ConsoleRuntime"/> has to be in the context of the running PowerShell thread, for example,
+        /// the callback of <see cref="PredictorInitializer.OnImport"/>.
+        /// The callbacks of <see cref="AzPredictor"/> are on a thread pool and it must not be used there.
+        /// </remarks>
+        PowerShell ConsoleRuntime { get; }
+
+        /// <summary>
+        /// Executes the PowerShell cmdlet in the current powershell session.
+        /// </summary>
+        IList<T> ExecuteScript<T>(string contents);
+    }
+}

tools/Az.Tools.Predictor/Az.Tools.Predictor/Utilities/PowerShellRunspaceUtilities.cs

@@ -0,0 +1,39 @@
+// ----------------------------------------------------------------------------------
+//
+// 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.Management.Automation.Runspaces;
+
+namespace Microsoft.Azure.PowerShell.Tools.AzPredictor.Utilities
+{
+    /// <summary>
+    /// A provider to create runspace in PowerShell.
+    /// </summary>
+    internal class PowerShellRunspaceUtilities
+    {
+        /// <summary>
+        /// Gets a new minimal runspace.
+        /// </summary>
+        public static Runspace GetMinimalRunspace()
+        {
+            // Create a mini runspace by remove the types and formats
+            InitialSessionState minimalState = InitialSessionState.CreateDefault2();
+            // Refer to the remarks for the property DefaultRunspace.
+            minimalState.Types.Clear();
+            minimalState.Formats.Clear();
+            var runspace = RunspaceFactory.CreateRunspace(minimalState);
+            runspace.Open();
+            return runspace;
+        }
+    }
+}

tools/Az.Tools.Predictor/Az.Tools.Predictor/Utilities/PowerShellRuntime.cs

@@ -27,7 +27,7 @@ namespace Microsoft.Azure.PowerShell.Tools.AzPredictor.Utilities
     /// <summary>
     /// A PowerShell environment to run PowerShell cmdlets and scripts.
     /// </summary>
-    internal class PowerShellRuntime : IDisposable
+    internal class PowerShellRuntime : IPowerShellRuntime, IDisposable
     {
         private PowerShell _runtime;
         private PowerShell Runtime
@@ -43,37 +43,13 @@ private PowerShell Runtime
             }
         }
 
-        private readonly Lazy<Runspace> _defaultRunspace = new(() =>
-                {
-                    // Create a mini runspace by remove the types and formats
-                    InitialSessionState minimalState = InitialSessionState.CreateDefault2();
-                    // Refer to the remarks for the property DefaultRunspace.
-                    minimalState.Types.Clear();
-                    minimalState.Formats.Clear();
-                    var runspace = RunspaceFactory.CreateRunspace(minimalState);
-                    runspace.Open();
-                    return runspace;
-                });
+        private readonly Lazy<Runspace> _defaultRunspace = new(() => PowerShellRunspaceUtilities.GetMinimalRunspace());
 
         /// <inheritdoc />
-        /// <remarks>
-        /// We don't pre-load Az service modules since they may not always be installed.
-        /// Creating the instance is at the first time this is called.
-        /// It can be slow. So the first call must not be in the path of the user interaction.
-        /// Loading too many modules can also impact user experience because that may add to much memory pressure at the same
-        /// time.
-        /// </remarks>
         public Runspace DefaultRunspace => _defaultRunspace.Value;
 
-        /// <summary>
-        /// The PowerShell environment that the module is imported into.
-        /// </summary>
-        /// <remarks>
-        /// The usage of <see cref="ConsoleRuntime"/> has to be in the context of the running PowerShell thread, for example,
-        /// the callback of <see cref="PredictorInitializer.OnImport"/>.
-        /// The callbacks of <see cref="AzPredictor"/> are on a thread pool and it must not be used there.
-        /// </remarks>
-        internal PowerShell ConsoleRuntime = PowerShell.Create(System.Management.Automation.RunspaceMode.CurrentRunspace);
+        /// <inheritdoc />
+        public PowerShell ConsoleRuntime { get; } = PowerShell.Create(System.Management.Automation.RunspaceMode.CurrentRunspace);
 
         public void Dispose()
         {
@@ -94,9 +70,7 @@ public void Dispose()
             }
         }
 
-        /// <summary>
-        /// Executes the PowerShell cmdlet in the current powershell session.
-        /// </summary>
+        /// <inheritdoc />
         public IList<T> ExecuteScript<T>(string contents)
         {
             Runtime.Commands.Clear();