Merge pull request #143 from PowerShell/daviwil/output-test-fix

Fix timing issues with OutputReader in host tests

Author: daviwil

test/PowerShellEditorServices.Test.Host/LanguageServerTests.cs

@@ -538,16 +538,31 @@ await requestContext.SendResult(
                     ResponseText = "John"
                 });
 
-            // Skip the initial script lines (4 script lines plus 3 blank lines)
-            await outputReader.ReadLines(7);
-            
-            // Verify output
-            string[] outputLines = await outputReader.ReadLines(5);
-            Assert.Equal("Name: John", outputLines[0]);
-            Assert.Equal("", outputLines[1]);
-            Assert.Equal("Key  Value", outputLines[2]);
-            Assert.Equal("---  -----", outputLines[3]);
-            Assert.Equal("Name John ", outputLines[4]);
+            // Skip the initial script lines (4 script lines plus 2 blank lines)
+            string[] scriptLines = await outputReader.ReadLines(6);
+
+            // In some cases an extra newline appears after the script lines.
+            // I have no idea why this happens, but it normally seems to occur
+            // on my local machine and not the CI server.  For now, adjust for
+            // it here.
+            string outputLine = await outputReader.ReadLine();
+            if (string.IsNullOrEmpty(outputLine))
+            {
+                outputLine = await outputReader.ReadLine();
+            }
+
+            // Verify the first line
+            Assert.Equal("Name: John", outputLine);
+
+            // Verify the rest of the output
+            string[] outputLines = await outputReader.ReadLines(4);
+            Assert.Equal("", outputLines[0]);
+            Assert.Equal("Key  Value", outputLines[1]);
+            Assert.Equal("---  -----", outputLines[2]);
+            Assert.Equal("Name John ", outputLines[3]);
+
+            // Wait for execution to complete
+            await evaluateTask;
         }
 
         private async Task SendOpenFileEvent(string filePath, bool waitForDiagnostics = true)

test/PowerShellEditorServices.Test.Host/OutputReader.cs

@@ -10,6 +10,7 @@
 using System.Collections.Generic;
 using System.Linq;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 using Xunit;
 
@@ -20,7 +21,7 @@ internal class OutputReader
         private AsyncQueue<OutputEventBody> outputQueue = new AsyncQueue<OutputEventBody>();
 
         private string currentOutputCategory;
-        private Queue<string> bufferedOutput = new Queue<string>();
+        private Queue<Tuple<string, bool>> bufferedOutput = new Queue<Tuple<string, bool>>();
 
         public OutputReader(ProtocolEndpoint protocolClient)
         {
@@ -31,34 +32,80 @@ public OutputReader(ProtocolEndpoint protocolClient)
 
         public async Task<string> ReadLine(string expectedOutputCategory = "stdout")
         {
-            if (this.bufferedOutput.Count > 0)
+            try
             {
-                Assert.Equal(expectedOutputCategory, this.currentOutputCategory);
+                bool lineHasNewLine = false;
+                string[] outputLines = null;
+                string nextOutputString = string.Empty;
 
-                return this.bufferedOutput.Dequeue();
-            }
-
-            // Execution reaches this point if a buffered line wasn't available
-            OutputEventBody nextOutputEvent = await this.outputQueue.DequeueAsync();
+                // Wait no longer than 7 seconds for output to come back
+                CancellationTokenSource cancellationSource = new CancellationTokenSource(7000);
 
-            Assert.Equal(expectedOutputCategory, nextOutputEvent.Category);
-            this.currentOutputCategory = nextOutputEvent.Category;
+                // Any lines in the buffer?
+                if (this.bufferedOutput.Count > 0)
+                {
+                    Assert.Equal(expectedOutputCategory, this.currentOutputCategory);
 
-            string[] outputLines = 
-                nextOutputEvent.Output.Split(
-                    new string[] { "\n", "\r\n" },
-                    StringSplitOptions.None);
+                    // Return the first buffered line
+                    var lineTuple = this.bufferedOutput.Dequeue();
+                    nextOutputString = lineTuple.Item1;
+                    lineHasNewLine = lineTuple.Item2;
+                }
 
-            // Buffer remaining lines
-            if (outputLines.Length > 1)
-            {
-                for (int i = 1; i < outputLines.Length; i++)
+                // Loop until we get a full line of output
+                while (!lineHasNewLine)
                 {
-                    this.bufferedOutput.Enqueue(outputLines[i]);
+                    // Execution reaches this point if a buffered line wasn't available
+                    Task<OutputEventBody> outputTask =
+                        this.outputQueue.DequeueAsync(
+                            cancellationSource.Token);
+                    OutputEventBody nextOutputEvent = await outputTask;
+
+                    // Verify that the output is of the expected type
+                    Assert.Equal(expectedOutputCategory, nextOutputEvent.Category);
+                    this.currentOutputCategory = nextOutputEvent.Category;
+
+                    // Split up the output into multiple lines
+                    outputLines = 
+                        nextOutputEvent.Output.Split(
+                            new string[] { "\n", "\r\n" },
+                            StringSplitOptions.None);
+
+                    // Add the first bit of output to the existing string
+                    nextOutputString += outputLines[0];
+
+                    // Have we found a newline now?
+                    lineHasNewLine =
+                        outputLines.Length > 1 ||
+                        nextOutputEvent.Output.EndsWith("\n");
+
+                    // Buffer any remaining lines for future reads
+                    if (outputLines.Length > 1)
+                    {
+                        for (int i = 1; i < outputLines.Length; i++)
+                        {
+                            this.bufferedOutput.Enqueue(
+                                new Tuple<string, bool>(
+                                    outputLines[i],
+
+                                    // The line has a newline if it's not the last segment or
+                                    // if the current output string ends with a newline
+                                    i < outputLines.Length - 1 ||
+                                    nextOutputEvent.Output.EndsWith("\n")));
+                        }
+                    }
+
+                    // At this point, the state of lineHasNewLine will determine
+                    // whether the loop continues to wait for another output
+                    // event that completes the current line.
                 }
-            }
 
-            return outputLines[0];
+                return nextOutputString;
+            }
+            catch (TaskCanceledException)
+            {
+                throw new TimeoutException("Timed out waiting for an input line.");
+            }
         }
 
         public async Task<string[]> ReadLines(int lineCount, string expectedOutputCategory = "stdout")