From 81f8b1745126f83d6183fe92aaad56487cc7e63a Mon Sep 17 00:00:00 2001
From: Tyrrrz <1935960+Tyrrrz@users.noreply.github.com>
Date: Tue, 16 May 2023 02:43:47 +0300
Subject: [PATCH] Fix cancellation tests

---
 .../Commands/CancellationTestCommand.cs       |  6 ++-
 CliFx.Tests/CancellationSpecs.cs              | 38 +++++++++++++++----
 2 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/CliFx.Tests.Dummy/Commands/CancellationTestCommand.cs b/CliFx.Tests.Dummy/Commands/CancellationTestCommand.cs
index 15964e3..96987b5 100644
--- a/CliFx.Tests.Dummy/Commands/CancellationTestCommand.cs
+++ b/CliFx.Tests.Dummy/Commands/CancellationTestCommand.cs
@@ -12,16 +12,18 @@ public class CancellationTestCommand : ICommand
     {
         try
         {
+            console.Output.WriteLine("Started.");
+
             await Task.Delay(
                 TimeSpan.FromSeconds(3),
                 console.RegisterCancellationHandler()
             );
 
-            console.Output.WriteLine("Completed successfully");
+            console.Output.WriteLine("Completed.");
         }
         catch (OperationCanceledException)
         {
-            console.Output.WriteLine("Cancelled");
+            console.Output.WriteLine("Cancelled.");
             throw;
         }
     }
diff --git a/CliFx.Tests/CancellationSpecs.cs b/CliFx.Tests/CancellationSpecs.cs
index 56c4e41..e6cd25f 100644
--- a/CliFx.Tests/CancellationSpecs.cs
+++ b/CliFx.Tests/CancellationSpecs.cs
@@ -4,6 +4,7 @@ using System.Text;
 using System.Threading;
 using System.Threading.Tasks;
 using CliFx.Tests.Utils;
+using CliFx.Tests.Utils.Extensions;
 using CliWrap;
 using FluentAssertions;
 using Xunit;
@@ -22,16 +23,29 @@ public class CancellationSpecs : SpecsBase
     public async Task I_can_configure_the_command_to_listen_to_the_interrupt_signal()
     {
         // Arrange
+        using var cts = new CancellationTokenSource();
+
+        // We need to send the cancellation request right after the process has registered
+        // a handler for the interrupt signal, otherwise the default handler will trigger
+        // and just kill the process.
+        void HandleStdOut(string line)
+        {
+            if (string.Equals(line, "Started.", StringComparison.OrdinalIgnoreCase))
+                cts.CancelAfter(TimeSpan.FromSeconds(0.2));
+        }
+
         var stdOutBuffer = new StringBuilder();
 
+        var pipeTarget = PipeTarget.Merge(
+            PipeTarget.ToDelegate(HandleStdOut),
+            PipeTarget.ToStringBuilder(stdOutBuffer)
+        );
+
         var command = Cli.Wrap("dotnet")
             .WithArguments(a => a
                 .Add(Dummy.Program.Location)
                 .Add("cancel-test")
-            ) | stdOutBuffer;
-
-        using var cts = new CancellationTokenSource();
-        cts.CancelAfter(TimeSpan.FromSeconds(0.2));
+            ) | pipeTarget;
 
         // Act & assert
         await Assert.ThrowsAnyAsync<OperationCanceledException>(async () =>
@@ -43,7 +57,10 @@ public class CancellationSpecs : SpecsBase
             )
         );
 
-        stdOutBuffer.ToString().Trim().Should().Be("Cancelled");
+        stdOutBuffer.ToString().Trim().Should().ConsistOfLines(
+            "Started.",
+            "Cancelled."
+        );
     }
 
     [Fact]
@@ -60,16 +77,18 @@ public class CancellationSpecs : SpecsBase
                 {
                     try
                     {
+                        console.Output.WriteLine("Started.");
+
                         await Task.Delay(
                             TimeSpan.FromSeconds(3),
                             console.RegisterCancellationHandler()
                         );
 
-                        console.Output.WriteLine("Completed successfully");
+                        console.Output.WriteLine("Completed.");
                     }
                     catch (OperationCanceledException)
                     {
-                        console.Output.WriteLine("Cancelled");
+                        console.Output.WriteLine("Cancelled.");
                         throw;
                     }
                 }
@@ -94,6 +113,9 @@ public class CancellationSpecs : SpecsBase
         exitCode.Should().NotBe(0);
 
         var stdOut = FakeConsole.ReadOutputString();
-        stdOut.Trim().Should().Be("Cancelled");
+        stdOut.Trim().Should().ConsistOfLines(
+            "Started.",
+            "Cancelled."
+        );
     }
 }
\ No newline at end of file