MegaBites-AI/Windows-powershell

1// Copyright (c) Microsoft Corporation.2// Licensed under the MIT License.3 4using System;5using System.Collections.Generic;6using System.Management.Automation.Language;7using System.Management.Automation.Subsystem;8using System.Management.Automation.Subsystem.Prediction;9using System.Threading;10using Xunit;11 12namespace PSTests.Sequential13{14    public class MyPredictor : ICommandPredictor15    {16        private readonly Guid _id;17        private readonly string _name, _description;18        private readonly bool _delay;19 20        public List<string> History { get; }21 22        public List<string> Results { get; }23 24        public List<string> AcceptedSuggestions { get; }25 26        public List<string> DisplayedSuggestions { get; }27 28        public static readonly MyPredictor SlowPredictor, FastPredictor;29 30        static MyPredictor()31        {32            SlowPredictor = new MyPredictor(33                Guid.NewGuid(),34                "Test Predictor #1",35                "Description for #1 predictor.",36                delay: true);37 38            FastPredictor = new MyPredictor(39                Guid.NewGuid(),40                "Test Predictor #2",41                "Description for #2 predictor.",42                delay: false);43        }44 45        private MyPredictor(Guid id, string name, string description, bool delay)46        {47            _id = id;48            _name = name;49            _description = description;50            _delay = delay;51 52            History = new List<string>();53            Results = new List<string>();54            AcceptedSuggestions = new List<string>();55            DisplayedSuggestions = new List<string>();56        }57 58        public void Clear()59        {60            History.Clear();61            Results.Clear();62            AcceptedSuggestions.Clear();63            DisplayedSuggestions.Clear();64        }65 66        #region "Interface implementation"67 68        public Guid Id => _id;69 70        public string Name => _name;71 72        public string Description => _description;73 74        bool ICommandPredictor.CanAcceptFeedback(PredictionClient client, PredictorFeedbackKind feedback) => true;75 76        public SuggestionPackage GetSuggestion(PredictionClient client, PredictionContext context, CancellationToken cancellationToken)77        {78            if (_delay)79            {80                // The delay is exaggerated to make the test reliable.81                // xUnit must spin up a lot tasks, which makes the test unreliable when the time difference between 'delay' and 'timeout' is small.82                Thread.Sleep(3000);83            }84 85            // You can get the user input from the AST.86            var userInput = context.InputAst.Extent.Text;87            var entries = new List<PredictiveSuggestion>88            {89                new PredictiveSuggestion($"'{userInput}' from '{client.Name}' - TEST-1 from {Name}"),90                new PredictiveSuggestion($"'{userInput}' from '{client.Name}' - TeSt-2 from {Name}"),91            };92 93            return new SuggestionPackage(56, entries);94        }95 96        public void OnSuggestionDisplayed(PredictionClient client, uint session, int countOrIndex)97        {98            DisplayedSuggestions.Add($"{client.Name}-{session}-{countOrIndex}");99        }100 101        public void OnSuggestionAccepted(PredictionClient client, uint session, string acceptedSuggestion)102        {103            AcceptedSuggestions.Add($"{client.Name}-{session}-{acceptedSuggestion}");104        }105 106        public void OnCommandLineAccepted(PredictionClient client, IReadOnlyList<string> history)107        {108            foreach (string item in history)109            {110                History.Add($"{client.Name}-{item}");111            }112        }113 114        public void OnCommandLineExecuted(PredictionClient client, string commandLine, bool success)115        {116            Results.Add($"{client.Name}-{commandLine}-{success}");117        }118 119        #endregion120    }121 122    public static class CommandPredictionTests123    {124        private const string Client = "PredictionTest";125        private const uint Session = 56;126        private static readonly PredictionClient predClient = new(Client, PredictionClientKind.Terminal);127 128        [Fact]129        public static void PredictInput()130        {131            const string Input = "Hello world";132            MyPredictor slow = MyPredictor.SlowPredictor;133            MyPredictor fast = MyPredictor.FastPredictor;134            Ast ast = Parser.ParseInput(Input, out Token[] tokens, out _);135 136            // Returns null when no predictor implementation registered137            List<PredictionResult> results = CommandPrediction.PredictInputAsync(predClient, ast, tokens).Result;138            Assert.Null(results);139 140            try141            {142                // Register 2 predictor implementations143                SubsystemManager.RegisterSubsystem<ICommandPredictor, MyPredictor>(slow);144                SubsystemManager.RegisterSubsystem(SubsystemKind.CommandPredictor, fast);145 146                // Expect the results from 'fast' predictor only b/c the 'slow' one147                // cannot finish before the specified timeout.148                // The specified timeout is exaggerated to make the test reliable.149                // xUnit must spin up a lot tasks, which makes the test unreliable when the time difference between 'delay' and 'timeout' is small.150                results = CommandPrediction.PredictInputAsync(predClient, ast, tokens, millisecondsTimeout: 1500).Result;151                Assert.Single(results);152 153                PredictionResult res = results[0];154                Assert.Equal(fast.Id, res.Id);155                Assert.Equal(Session, res.Session);156                Assert.Equal(2, res.Suggestions.Count);157                Assert.Equal($"'{Input}' from '{Client}' - TEST-1 from {fast.Name}", res.Suggestions[0].SuggestionText);158                Assert.Equal($"'{Input}' from '{Client}' - TeSt-2 from {fast.Name}", res.Suggestions[1].SuggestionText);159 160                // Expect the results from both 'slow' and 'fast' predictors161                // Same here -- the specified timeout is exaggerated to make the test reliable.162                // xUnit must spin up a lot tasks, which makes the test unreliable when the time difference between 'delay' and 'timeout' is small.163                results = CommandPrediction.PredictInputAsync(predClient, ast, tokens, millisecondsTimeout: 4000).Result;164                Assert.Equal(2, results.Count);165 166                PredictionResult res1 = results[0];167                Assert.Equal(slow.Id, res1.Id);168                Assert.Equal(Session, res1.Session);169                Assert.Equal(2, res1.Suggestions.Count);170                Assert.Equal($"'{Input}' from '{Client}' - TEST-1 from {slow.Name}", res1.Suggestions[0].SuggestionText);171                Assert.Equal($"'{Input}' from '{Client}' - TeSt-2 from {slow.Name}", res1.Suggestions[1].SuggestionText);172 173                PredictionResult res2 = results[1];174                Assert.Equal(fast.Id, res2.Id);175                Assert.Equal(Session, res2.Session);176                Assert.Equal(2, res2.Suggestions.Count);177                Assert.Equal($"'{Input}' from '{Client}' - TEST-1 from {fast.Name}", res2.Suggestions[0].SuggestionText);178                Assert.Equal($"'{Input}' from '{Client}' - TeSt-2 from {fast.Name}", res2.Suggestions[1].SuggestionText);179            }180            finally181            {182                SubsystemManager.UnregisterSubsystem<ICommandPredictor>(slow.Id);183                SubsystemManager.UnregisterSubsystem(SubsystemKind.CommandPredictor, fast.Id);184            }185        }186 187        [Fact]188        public static void Feedback()189        {190            MyPredictor slow = MyPredictor.SlowPredictor;191            MyPredictor fast = MyPredictor.FastPredictor;192 193            slow.Clear();194            fast.Clear();195 196            try197            {198                // Register 2 predictor implementations199                SubsystemManager.RegisterSubsystem<ICommandPredictor, MyPredictor>(slow);200                SubsystemManager.RegisterSubsystem(SubsystemKind.CommandPredictor, fast);201 202                var history = new[] { "hello", "world" };203                var ids = new HashSet<Guid> { slow.Id, fast.Id };204 205                CommandPrediction.OnCommandLineAccepted(predClient, history);206                CommandPrediction.OnCommandLineExecuted(predClient, "last_input", true);207                CommandPrediction.OnSuggestionDisplayed(predClient, slow.Id, Session, 2);208                CommandPrediction.OnSuggestionDisplayed(predClient, fast.Id, Session, -1);209                CommandPrediction.OnSuggestionAccepted(predClient, slow.Id, Session, "Yeah");210 211                // The feedback calls are queued in thread pool, so let's wait a bit to make sure the calls are done.212                while (slow.History.Count == 0 || fast.History.Count == 0 ||213                       slow.Results.Count == 0 || fast.Results.Count == 0 ||214                       slow.DisplayedSuggestions.Count == 0 || fast.DisplayedSuggestions.Count == 0 ||215                       slow.AcceptedSuggestions.Count == 0)216                {217                    Thread.Sleep(300);218                }219 220                Assert.Equal(2, slow.History.Count);221                Assert.Equal($"{Client}-{history[0]}", slow.History[0]);222                Assert.Equal($"{Client}-{history[1]}", slow.History[1]);223 224                Assert.Equal(2, fast.History.Count);225                Assert.Equal($"{Client}-{history[0]}", fast.History[0]);226                Assert.Equal($"{Client}-{history[1]}", fast.History[1]);227 228                Assert.Single(slow.Results);229                Assert.Equal($"{Client}-last_input-True", slow.Results[0]);230 231                Assert.Single(fast.Results);232                Assert.Equal($"{Client}-last_input-True", fast.Results[0]);233 234                Assert.Single(slow.DisplayedSuggestions);235                Assert.Equal($"{Client}-{Session}-2", slow.DisplayedSuggestions[0]);236 237                Assert.Single(fast.DisplayedSuggestions);238                Assert.Equal($"{Client}-{Session}--1", fast.DisplayedSuggestions[0]);239 240                Assert.Single(slow.AcceptedSuggestions);241                Assert.Equal($"{Client}-{Session}-Yeah", slow.AcceptedSuggestions[0]);242 243                Assert.Empty(fast.AcceptedSuggestions);244            }245            finally246            {247                SubsystemManager.UnregisterSubsystem<ICommandPredictor>(slow.Id);248                SubsystemManager.UnregisterSubsystem(SubsystemKind.CommandPredictor, fast.Id);249            }250        }251    }252}253