Set discovery batch size to 1000 (#3896)

* Set discovery batch size to 1000

* Address review comments

* Fix test

Co-authored-by: Amaury Levé <[email protected]>

src/Microsoft.TestPlatform.Utilities/InferRunSettingsHelper.cs

@@ -24,6 +24,7 @@ namespace Microsoft.VisualStudio.TestPlatform.Utilities;
 public class InferRunSettingsHelper
 {
     private const string DesignModeNodeName = "DesignMode";
+    private const string BatchSizeNodeName = "BatchSize";
     private const string CollectSourceInformationNodeName = "CollectSourceInformation";
     private const string RunSettingsNodeName = "RunSettings";
     private const string RunConfigurationNodeName = "RunConfiguration";
@@ -33,6 +34,7 @@ public class InferRunSettingsHelper
     private const string TargetDevice = "TargetDevice";
 
     private const string DesignModeNodePath = @"/RunSettings/RunConfiguration/DesignMode";
+    private const string BatchSizeNodePath = @"/RunSettings/RunConfiguration/BatchSize";
     private const string CollectSourceInformationNodePath = @"/RunSettings/RunConfiguration/CollectSourceInformation";
     private const string RunConfigurationNodePath = @"/RunSettings/RunConfiguration";
     private const string TargetPlatformNodePath = @"/RunSettings/RunConfiguration/TargetPlatform";
@@ -201,6 +203,16 @@ public static void UpdateDesignMode(XmlDocument runSettingsDocument, bool design
         AddNodeIfNotPresent(runSettingsDocument, DesignModeNodePath, DesignModeNodeName, designModeValue);
     }
 
+    /// <summary>
+    /// Updates the <c>RunConfiguration.BatchSize</c> value for a run settings. Doesn't do anything if the value is already set.
+    /// </summary>
+    /// <param name="runSettingsDocument">Document for runsettings xml</param>
+    /// <param name="batchSizeValue">Value to set</param>
+    public static void UpdateBatchSize(XmlDocument runSettingsDocument, long batchSizeValue)
+    {
+        AddNodeIfNotPresent(runSettingsDocument, BatchSizeNodePath, BatchSizeNodeName, batchSizeValue);
+    }
+
     /// <summary>
     /// Updates the <c>RunConfiguration.CollectSourceInformation</c> value for a run settings. Doesn't do anything if the value is already set.
     /// </summary>

src/Microsoft.TestPlatform.Utilities/PublicAPI/PublicAPI.Shipped.txt

@@ -38,3 +38,5 @@ static Microsoft.VisualStudio.TestPlatform.Utilities.MSTestSettingsUtilities.Imp
 static Microsoft.VisualStudio.TestPlatform.Utilities.MSTestSettingsUtilities.IsLegacyTestSettingsFile(string? settingsFile) -> bool
 static Microsoft.VisualStudio.TestPlatform.Utilities.ParallelRunSettingsUtilities.UpdateRunSettingsWithParallelSettingIfNotConfigured(System.Xml.XPath.XPathNavigator! navigator) -> void
 static Microsoft.VisualStudio.TestPlatform.Utilities.StringExtensions.Tokenize(this string? input, char separator, char escape) -> System.Collections.Generic.IEnumerable<string!>!
+static Microsoft.VisualStudio.TestPlatform.Utilities.InferRunSettingsHelper.UpdateBatchSize(System.Xml.XmlDocument! runSettingsDocument, long batchSizeValue) -> void
+

src/vstest.console/CommandLine/CommandLineOptions.cs

@@ -24,9 +24,14 @@ namespace Microsoft.VisualStudio.TestPlatform.CommandLine;
 internal class CommandLineOptions
 {
     /// <summary>
-    /// The default batch size.
+    /// The default batch size for Run.
     /// </summary>
-    public const long DefaultBatchSize = 10;
+    public const long DefaultRunBatchSize = 10;
+
+    /// <summary>
+    /// The default batch size for Discovery.
+    /// </summary>
+    public const long DefaultDiscoveryBatchSize = 1000;
 
     /// <summary>
     /// The use vsix extensions key.
@@ -62,7 +67,7 @@ internal static CommandLineOptions Instance
     /// </summary>
     internal CommandLineOptions()
     {
-        BatchSize = DefaultBatchSize;
+        BatchSize = DefaultRunBatchSize;
         TestStatsEventTimeout = _defaultRetrievalTimeout;
         FileHelper = new FileHelper();
         FilePatternParser = new FilePatternParser();

src/vstest.console/TestPlatformHelpers/TestRequestManager.cs

@@ -174,6 +174,7 @@ public void DiscoverTests(
                 runsettings,
                 discoveryPayload.Sources.ToList(),
                 discoveryEventsRegistrar,
+                isDiscovery: true,
                 out string updatedRunsettings,
                 out IDictionary<string, Architecture> sourceToArchitectureMap,
                 out IDictionary<string, Framework> sourceToFrameworkMap))
@@ -265,6 +266,7 @@ public void RunTests(
         ProtocolConfig protocolConfig)
     {
         EqtTrace.Info("TestRequestManager.RunTests: run tests started.");
+        testRunRequestPayload.RunSettings ??= "<RunSettings></RunSettings>";
         var runsettings = testRunRequestPayload.RunSettings;
 
         if (testRunRequestPayload.TestPlatformOptions != null)
@@ -281,6 +283,7 @@ public void RunTests(
                 runsettings!,
                 sources!,
                 testRunEventsRegistrar,
+                isDiscovery: false,
                 out string updatedRunsettings,
                 out IDictionary<string, Architecture> sourceToArchitectureMap,
                 out IDictionary<string, Framework> sourceToFrameworkMap))
@@ -468,7 +471,8 @@ public void StartTestSession(
         if (UpdateRunSettingsIfRequired(
                 payload.RunSettings,
                 payload.Sources,
-                null,
+                registrar: null,
+                isDiscovery: false,
                 out string updatedRunsettings,
                 out IDictionary<string, Architecture> sourceToArchitectureMap,
                 out IDictionary<string, Framework> sourceToFrameworkMap))
@@ -652,6 +656,7 @@ private bool UpdateRunSettingsIfRequired(
         string runsettingsXml,
         IList<string>? sources,
         IBaseTestEventsRegistrar? registrar,
+        bool isDiscovery,
         out string updatedRunSettingsXml,
         out IDictionary<string, Architecture> sourceToArchitectureMap,
         out IDictionary<string, Framework> sourceToFrameworkMap)
@@ -800,6 +805,7 @@ private bool UpdateRunSettingsIfRequired(
         settingsUpdated |= UpdateCollectSourceInformation(document, runConfiguration);
         settingsUpdated |= UpdateTargetDevice(navigator, document);
         settingsUpdated |= AddOrUpdateConsoleLogger(document, runConfiguration, loggerRunSettings);
+        settingsUpdated |= AddOrUpdateBatchSize(document, runConfiguration, isDiscovery);
 
         updatedRunSettingsXml = navigator.OuterXml;
 
@@ -909,6 +915,27 @@ private bool UpdateDesignMode(XmlDocument document, RunConfiguration runConfigur
         return updateRequired;
     }
 
+    internal /* for testing purposes */ static bool AddOrUpdateBatchSize(XmlDocument document, RunConfiguration runConfiguration, bool isDiscovery)
+    {
+        // On run keep it as is to fall back to the current default value (which is 10 right now).
+        if (!isDiscovery)
+        {
+            // We did not update runnsettings.
+            return false;
+        }
+
+        // If user is already setting batch size via runsettings or CLI args; we skip.
+        bool updateRequired = !runConfiguration.BatchSizeSet;
+        if (updateRequired)
+        {
+            InferRunSettingsHelper.UpdateBatchSize(
+                document,
+                CommandLineOptions.DefaultDiscoveryBatchSize);
+        }
+
+        return updateRequired;
+    }
+
     private static void CheckSourcesForCompatibility(
         Framework chosenFramework,
         Architecture chosenPlatform,

test/vstest.console.UnitTests/CommandLine/CommandLineOptionsTests.cs

@@ -39,6 +39,12 @@ public void CommandLineOptionsDefaultBatchSizeIsTen()
         Assert.AreEqual(10, CommandLineOptions.Instance.BatchSize);
     }
 
+    [TestMethod]
+    public void CommandLineOptionsDiscoveryDefaultBatchSizeIsThousand()
+    {
+        Assert.AreEqual(1000, CommandLineOptions.DefaultDiscoveryBatchSize);
+    }
+
     [TestMethod]
     public void CommandLineOptionsDefaultTestRunStatsEventTimeoutIsOnePointFiveSec()
     {

test/vstest.console.UnitTests/TestPlatformHelpers/TestRequestManagerTests.cs

@@ -8,6 +8,7 @@
 using System.Runtime.Versioning;
 using System.Threading;
 using System.Threading.Tasks;
+using System.Xml;
 
 using FluentAssertions;
 
@@ -219,8 +220,8 @@ public void DiscoverTestsShouldCallTestPlatformAndSucceed()
         Assert.AreEqual("a", actualDiscoveryCriteria.Sources.First(), "First Source in list is incorrect");
         Assert.AreEqual("b", actualDiscoveryCriteria.Sources.ElementAt(1), "Second Source in list is incorrect");
 
-        // Default frequency is set to 10, unless specified in runsettings.
-        Assert.AreEqual(10, actualDiscoveryCriteria.FrequencyOfDiscoveredTestsEvent);
+        // Default frequency is set to BatchSize (which is set to 1000).
+        Assert.AreEqual(1000, actualDiscoveryCriteria.FrequencyOfDiscoveredTestsEvent);
 
         mockDiscoveryRegistrar.Verify(md => md.RegisterDiscoveryEvents(It.IsAny<IDiscoveryRequest>()), Times.Once);
         mockDiscoveryRegistrar.Verify(md => md.UnregisterDiscoveryEvents(It.IsAny<IDiscoveryRequest>()), Times.Once);
@@ -2518,6 +2519,79 @@ public void StopTestSessionShouldPropagateExceptionWhenKillSessionThrows()
             Times.Once);
     }
 
+    [TestMethod]
+    public void AddOrUpdateBatchSizeWhenNotDiscoveryReturnsFalseAndDoesNotUpdateXmlDocument()
+    {
+        // Arrange
+        var xmlDocument = new XmlDocument();
+        var configuration = new RunConfiguration();
+
+        // Act
+        var result = TestRequestManager.AddOrUpdateBatchSize(xmlDocument, configuration, false);
+
+        // Assert
+        Assert.IsFalse(result);
+        Assert.AreEqual("", xmlDocument.OuterXml);
+    }
+
+    [TestMethod]
+    public void AddOrUpdateBatchSizeWhenBatchSizeSetReturnsFalse()
+    {
+        // Arrange
+        var xmlDocument = new XmlDocument();
+        var configuration = new RunConfiguration { BatchSize = 10 };
+
+        // Sanity check
+        Assert.IsTrue(configuration.BatchSizeSet);
+
+        // Act
+        var result = TestRequestManager.AddOrUpdateBatchSize(xmlDocument, configuration, true);
+
+        // Assert
+        Assert.IsFalse(result);
+        Assert.AreEqual("", xmlDocument.OuterXml);
+    }
+
+    [TestMethod]
+    public void AddOrUpdateBatchSizeSetsBatchSize()
+    {
+        // Arrange
+        var xmlDocument = new XmlDocument();
+        xmlDocument.LoadXml("""
+            <RunSettings>
+                <RunConfiguration>
+                </RunConfiguration>
+            </RunSettings>
+            """);
+        var configuration = new RunConfiguration();
+
+        // Act
+        var result = TestRequestManager.AddOrUpdateBatchSize(xmlDocument, configuration, true);
+
+        // Assert
+        Assert.IsTrue(result);
+        Assert.AreEqual("<RunSettings><RunConfiguration><BatchSize>1000</BatchSize></RunConfiguration></RunSettings>", xmlDocument.OuterXml);
+    }
+
+    [TestMethod]
+    public void AddOrUpdateBatchSizeSetsRunConfigurationAndBatchSize()
+    {
+        // Arrange
+        var xmlDocument = new XmlDocument();
+        xmlDocument.LoadXml("""
+            <RunSettings>
+            </RunSettings>
+            """);
+        var configuration = new RunConfiguration();
+
+        // Act
+        var result = TestRequestManager.AddOrUpdateBatchSize(xmlDocument, configuration, true);
+
+        // Assert
+        Assert.IsTrue(result);
+        Assert.AreEqual("<RunSettings><RunConfiguration><BatchSize>1000</BatchSize></RunConfiguration></RunSettings>", xmlDocument.OuterXml);
+    }
+
     private static DiscoveryRequestPayload CreateDiscoveryPayload(string runsettings)
     {
         var discoveryPayload = new DiscoveryRequestPayload