From 82780ead28d0784b9f86805d88faa5f4cd4bcf15 Mon Sep 17 00:00:00 2001
From: Sanan Yuzbashiyev <Sanan07@users.noreply.github.com>
Date: Tue, 11 May 2021 21:28:49 +0200
Subject: [PATCH] Aborting discovery (#2896)

---
 .../ParallelDiscoveryEventsHandler.cs         |  9 ++++++++
 .../Parallel/ParallelProxyDiscoveryManager.cs | 13 +++++++++--
 .../ParallelProxyDiscoveryManagerTests.cs     | 22 +++++++++++++++++++
 3 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/src/Microsoft.TestPlatform.CrossPlatEngine/Client/Parallel/ParallelDiscoveryEventsHandler.cs b/src/Microsoft.TestPlatform.CrossPlatEngine/Client/Parallel/ParallelDiscoveryEventsHandler.cs
index 93d1115754..1552ca924c 100644
--- a/src/Microsoft.TestPlatform.CrossPlatEngine/Client/Parallel/ParallelDiscoveryEventsHandler.cs
+++ b/src/Microsoft.TestPlatform.CrossPlatEngine/Client/Parallel/ParallelDiscoveryEventsHandler.cs
@@ -14,6 +14,8 @@ namespace Microsoft.VisualStudio.TestPlatform.CrossPlatEngine.Client.Parallel
     using Microsoft.VisualStudio.TestPlatform.ObjectModel.Engine;
     using Microsoft.VisualStudio.TestPlatform.ObjectModel.Logging;
 
+    using CommonResources = Common.Resources.Resources;
+
     /// <summary>
     /// ParallelDiscoveryEventsHandler for handling the discovery events in case of parallel discovery
     /// </summary>
@@ -122,6 +124,13 @@ public void HandleRawMessage(string rawMessage)
             // Always aggregate data, deserialize and raw for complete events
             var message = this.dataSerializer.DeserializeMessage(rawMessage);
 
+            // Do not send CancellationRequested message to Output window in IDE, as it is not useful for user
+            if (string.Equals(message.MessageType, MessageType.TestMessage)
+                && rawMessage.IndexOf(CommonResources.CancellationRequested) >= 0)
+            {
+                return;
+            }
+
             // Do not deserialize further
             if (!string.Equals(MessageType.DiscoveryComplete, message.MessageType))
             {
diff --git a/src/Microsoft.TestPlatform.CrossPlatEngine/Client/Parallel/ParallelProxyDiscoveryManager.cs b/src/Microsoft.TestPlatform.CrossPlatEngine/Client/Parallel/ParallelProxyDiscoveryManager.cs
index ec6a96acfb..6307ea41d3 100644
--- a/src/Microsoft.TestPlatform.CrossPlatEngine/Client/Parallel/ParallelProxyDiscoveryManager.cs
+++ b/src/Microsoft.TestPlatform.CrossPlatEngine/Client/Parallel/ParallelProxyDiscoveryManager.cs
@@ -38,6 +38,9 @@ internal class ParallelProxyDiscoveryManager : ParallelOperationManager<IProxyDi
 
         private IRequestData requestData;
 
+        // This field indicates if abort was requested by testplatform (user)
+        private bool discoveryAbortRequested = false;
+
         #endregion
 
         #region Concurrency Keeper Objects
@@ -89,6 +92,7 @@ public void DiscoverTests(DiscoveryCriteria discoveryCriteria, ITestDiscoveryEve
         /// <inheritdoc/>
         public void Abort()
         {
+            this.discoveryAbortRequested = true;
             this.DoActionOnAllManagers((proxyManager) => proxyManager.Abort(), doActionsInParallel: true);
         }
 
@@ -121,8 +125,13 @@ public bool HandlePartialDiscoveryComplete(IProxyDiscoveryManager proxyDiscovery
                 }
             }
 
-            // Discovery is completed. Schedule the clean up for managers and handlers.
-            if (allDiscoverersCompleted)
+            /*
+             If discovery is complete or discovery aborting was requsted by testPlatfrom(user)
+             we need to stop all ongoing discoveries, because we want to separate aborting request
+             when testhost crashed by itself and when user requested it (f.e. through TW)
+             Schedule the clean up for managers and handlers.
+            */
+            if (allDiscoverersCompleted || discoveryAbortRequested)
             {
                 // Reset enumerators
                 this.sourceEnumerator = null;
diff --git a/test/Microsoft.TestPlatform.CrossPlatEngine.UnitTests/Client/Parallel/ParallelProxyDiscoveryManagerTests.cs b/test/Microsoft.TestPlatform.CrossPlatEngine.UnitTests/Client/Parallel/ParallelProxyDiscoveryManagerTests.cs
index bcaee18ad0..d9f85b420b 100644
--- a/test/Microsoft.TestPlatform.CrossPlatEngine.UnitTests/Client/Parallel/ParallelProxyDiscoveryManagerTests.cs
+++ b/test/Microsoft.TestPlatform.CrossPlatEngine.UnitTests/Client/Parallel/ParallelProxyDiscoveryManagerTests.cs
@@ -122,6 +122,28 @@ public void DiscoveryTestsShouldProcessAllSourcesOnDiscoveryAbortsForAnySource()
             Assert.AreEqual(2, processedSources.Count, "All Sources must be processed.");
         }
 
+        /// <summary>
+        ///  Create ParallelProxyDiscoveryManager with parallel level 1 and two sources,
+        ///  Overall discovery should stop, if aborting was requested
+        /// </summary>
+        [TestMethod]
+        public void DiscoveryTestsShouldStopDiscoveryIfAbortionWasRequested()
+        {
+            // Since the hosts are aborted, total aggregated tests sent across will be -1
+            var discoveryManagerMock = new Mock<IProxyDiscoveryManager>();
+            this.createdMockManagers.Add(discoveryManagerMock);
+            var parallelDiscoveryManager = this.SetupDiscoveryManager(() => discoveryManagerMock.Object, 1, true, totalTests: -1);
+
+            Task.Run(() =>
+            {
+                parallelDiscoveryManager.DiscoverTests(this.testDiscoveryCriteria, this.mockHandler.Object);
+                parallelDiscoveryManager.Abort();
+            });
+
+            Assert.IsTrue(this.discoveryCompleted.Wait(taskTimeout), "Test discovery not completed.");
+            Assert.AreEqual(1, processedSources.Count, "One source should be processed.");
+        }
+
         [TestMethod]
         public void DiscoveryTestsShouldProcessAllSourceIfOneDiscoveryManagerIsStarved()
         {