Update code and add tests

src/Microsoft.TestPlatform.CrossPlatEngine/Client/Parallel/ParallelDiscoveryDataAggregator.cs

@@ -72,7 +72,7 @@ public IDictionary<string, object> GetAggregatedDiscoveryDataMetrics()
     /// Aggregate discovery data
     /// Must be thread-safe as this is expected to be called by parallel managers
     /// </summary>
-    public void Aggregate(DiscoveryCompleteEventArgs discoveryCompleteEventArgs)
+    public void Aggregate(DiscoveryCompleteEventArgs discoveryCompleteEventArgs!!)
     {
         lock (_dataUpdateSyncObject)
         {
@@ -93,7 +93,7 @@ public void Aggregate(DiscoveryCompleteEventArgs discoveryCompleteEventArgs)
         AggregateDiscoveryDataMetrics(discoveryCompleteEventArgs.Metrics);
     }
 
-    internal /* for testing purposes */ void AggregateDiscoveryStatus(
+    private void AggregateDiscoveryStatus(
         IEnumerable<string> notDiscoveredSources,
         IEnumerable<string> partiallyDiscoveredSources,
         IEnumerable<string> fullyDiscoveredSources)
@@ -143,7 +143,7 @@ public void Aggregate(DiscoveryCompleteEventArgs discoveryCompleteEventArgs)
     /// Aggregates the metrics from Test Host Process.
     /// </summary>
     /// <param name="metrics"></param>
-    internal /* for testing purposes */ void AggregateDiscoveryDataMetrics(IDictionary<string, object>? metrics)
+    private void AggregateDiscoveryDataMetrics(IDictionary<string, object>? metrics)
     {
         if (metrics == null || metrics.Count == 0 || _metricsAggregator == null)
         {

src/Microsoft.TestPlatform.CrossPlatEngine/Client/Parallel/ParallelDiscoveryEventsHandler.cs

@@ -62,10 +62,13 @@ public void HandleDiscoveryComplete(DiscoveryCompleteEventArgs discoveryComplete
         // Aggregate data for final discovery complete
         _discoveryDataAggregator.Aggregate(discoveryCompleteEventArgs);
 
-        // we get discovery complete events from each host process
-        // so we cannot "complete" the actual operation until all sources are consumed
-        // We should not block last chunk results while we aggregate overall discovery data
-        if (lastChunk?.Any() == true)
+        // We get DiscoveryComplete events from each ProxyDiscoveryManager (each testhost) and
+        // they contain the last chunk of tests that were discovered in that particular testhost.
+        // We want to send them to the IDE, but we don't want to tell it that discovery finished,
+        // because other sources are still being discovered. So we translate the last chunk to a
+        // normal TestsDiscovered event and send that to the IDE. Once all sources are completed,
+        // we will send a single DiscoveryComplete.
+        if (lastChunk != null)
         {
             ConvertToRawMessageAndSend(MessageType.TestCasesFound, lastChunk);
             HandleDiscoveredTests(lastChunk);
@@ -76,15 +79,16 @@ public void HandleDiscoveryComplete(DiscoveryCompleteEventArgs discoveryComplete
         var parallelDiscoveryComplete = _parallelProxyDiscoveryManager.HandlePartialDiscoveryComplete(
             _proxyDiscoveryManager,
             discoveryCompleteEventArgs.TotalCount,
-            null, // lastChunk should be null as we already sent this data above
+            null, // Set lastChunk to null, because we already sent the data using HandleDiscoveredTests above.
             discoveryCompleteEventArgs.IsAborted);
+
         if (!parallelDiscoveryComplete)
         {
             return;
         }
 
-        // As we immediately return results to IDE in case of aborting
-        // we need to set isAborted = true and totalTests = -1
+        // As we immediately return results to IDE in case of aborting we need to set
+        // isAborted = true and totalTests = -1
         if (_parallelProxyDiscoveryManager.IsAbortRequested)
         {
             _discoveryDataAggregator.Aggregate(new(-1, true));

src/Microsoft.TestPlatform.CrossPlatEngine/Client/ProxyDiscoveryManager.cs

@@ -235,12 +235,9 @@ public void Close()
     /// <inheritdoc/>
     public void HandleDiscoveryComplete(DiscoveryCompleteEventArgs discoveryCompleteEventArgs, IEnumerable<TestCase> lastChunk)
     {
-        if (lastChunk != null)
-        {
-            // When discovery is complete then the last discovered source is still marked
-            // as partially discovered, so we need to mark it as fully discovered.
-            _discoverySourceStatusCache.MarkTheLastChunkSourcesAsFullyDiscovered(lastChunk);
-        }
+        // When discovery is complete then the last discovered source is still marked
+        // as partially discovered, so we need to mark it as fully discovered.
+        _discoverySourceStatusCache.MarkTheLastChunkSourcesAsFullyDiscovered(lastChunk);
 
         _baseTestDiscoveryEventsHandler.HandleDiscoveryComplete(
             new(discoveryCompleteEventArgs.TotalCount, discoveryCompleteEventArgs.IsAborted)
@@ -249,8 +246,7 @@ public void HandleDiscoveryComplete(DiscoveryCompleteEventArgs discoveryComplete
                 NotDiscoveredSources = _discoverySourceStatusCache.GetSourcesWithStatus(DiscoveryStatus.NotDiscovered),
                 PartiallyDiscoveredSources = _discoverySourceStatusCache.GetSourcesWithStatus(DiscoveryStatus.PartiallyDiscovered),
                 FullyDiscoveredSources = _discoverySourceStatusCache.GetSourcesWithStatus(DiscoveryStatus.FullyDiscovered),
-            }
-            , lastChunk);
+            }, lastChunk);
     }
 
     /// <inheritdoc/>

src/Microsoft.TestPlatform.CrossPlatEngine/Discovery/DiscoveryManager.cs

@@ -39,7 +39,7 @@ public class DiscoveryManager : IDiscoveryManager
     private DiscoveryCriteria _discoveryCriteria;
     private readonly CancellationTokenSource _cancellationTokenSource = new();
     private readonly DiscoverySourceStatusCache _discoverySourceStatusCache = new();
-    
+
     /// <summary>
     /// Initializes a new instance of the <see cref="DiscoveryManager"/> class.
     /// </summary>

src/Microsoft.TestPlatform.CrossPlatEngine/Discovery/DiscoverySourceStatusCache.cs

@@ -16,7 +16,7 @@ internal class DiscoverySourceStatusCache
 
     private string? _previousSource;
 
-    public void MarkSourcesWithStatus(IEnumerable<string?> sources, DiscoveryStatus status)
+    public void MarkSourcesWithStatus(IEnumerable<string?>? sources, DiscoveryStatus status)
     {
         if (sources is null)
         {
@@ -40,27 +40,33 @@ public void MarkSourcesWithStatus(IEnumerable<string?> sources, DiscoveryStatus
 
                     return status;
                 },
-                (_, _) =>
+                (_, previousStatus) =>
                 {
+                    if (previousStatus == DiscoveryStatus.FullyDiscovered && status != DiscoveryStatus.FullyDiscovered
+                        || previousStatus == DiscoveryStatus.PartiallyDiscovered && status == DiscoveryStatus.NotDiscovered)
+                    {
+                        EqtTrace.Warning($"DiscoverySourceStatusCache.MarkSourcesWithStatus: Downgrading source status from {previousStatus} to {status}.");
+                    }
+
                     EqtTrace.Info($"DiscoverySourceStatusCache.MarkSourcesWithStatus: Marking {source} with {status} status.");
                     return status;
                 });
         }
     }
 
-    public void MarkTheLastChunkSourcesAsFullyDiscovered(IEnumerable<TestCase> lastChunk)
+    public void MarkTheLastChunkSourcesAsFullyDiscovered(IEnumerable<TestCase>? lastChunk)
     {
         // When all testcases count in project is dividable by chunk size (e.g. 100 tests and
         // chunk size of 10) then lastChunk is coming as empty. In this case, we need to take
         // the lastSource and mark it as FullyDiscovered.
-        IEnumerable<string?> lastChunkSources = lastChunk.Any()
+        IEnumerable<string?> lastChunkSources = lastChunk?.Any() == true
             ? lastChunk.Select<TestCase, string?>(testcase => testcase.Source)
             : new[] { _previousSource };
 
         MarkSourcesWithStatus(lastChunkSources, DiscoveryStatus.FullyDiscovered);
     }
 
-    public void MarkSourcesBasedOnDiscoveredTestCases(IEnumerable<TestCase> testCases)
+    public void MarkSourcesBasedOnDiscoveredTestCases(IEnumerable<TestCase>? testCases)
     {
         if (testCases is null)
         {
@@ -92,7 +98,7 @@ public List<string> GetSourcesWithStatus(DiscoveryStatus discoveryStatus)
         => GetSourcesWithStatus(discoveryStatus, _sourcesWithDiscoveryStatus);
 
     public static List<string> GetSourcesWithStatus(DiscoveryStatus discoveryStatus,
-        ConcurrentDictionary<string, DiscoveryStatus> sourcesWithDiscoveryStatus)
+        ConcurrentDictionary<string, DiscoveryStatus>? sourcesWithDiscoveryStatus)
     {
         // If by some accident SourcesWithDiscoveryStatus map is empty we will return empty list
         return sourcesWithDiscoveryStatus is null || sourcesWithDiscoveryStatus.IsEmpty

src/Microsoft.TestPlatform.ObjectModel/Client/Events/DiscoveryCompleteEventArgs.cs

@@ -27,11 +27,6 @@ public DiscoveryCompleteEventArgs(long totalTests, bool isAborted,
         IList<string> partiallyDiscoveredSources,
         IList<string> notDiscoveredSources)
     {
-        // This event is always raised from the client side, while the total count of tests is maintained
-        // only at the testhost end. In case of a discovery abort (various reasons including crash), it is
-        // not possible to get a list of total tests from testhost. Hence we enforce a -1 count.
-        Debug.Assert((!isAborted || -1 == totalTests), "If discovery request is aborted totalTest should be -1.");
-
         TotalCount = totalTests;
         IsAborted = isAborted;