Optimize TestMethodRunner: fix double enumeration, reduce allocations

src/Adapter/MSTestAdapter.PlatformServices/Discovery/AssemblyEnumerator.cs

@@ -269,34 +269,18 @@ private static bool TryUnfoldITestDataSources(UnitTestElement test, DiscoveryTes
     private static bool TryUnfoldITestDataSource(ITestDataSource dataSource, UnitTestElement test, ReflectionTestMethodInfo methodInfo, List<UnitTestElement> tests, ref int globalTestCaseIndex)
     {
         // Otherwise, unfold the data source and verify it can be serialized.
-        IEnumerable<object?[]>? data;
 
-        // This code is to discover tests. To run the tests code is in TestMethodRunner.ExecuteDataSourceBasedTests.
-        // Any change made here should be reflected in TestMethodRunner.ExecuteDataSourceBasedTests as well.
-        data = dataSource.GetData(methodInfo);
+        // This code is to discover tests. To run the tests code is in TestMethodRunner.TryExecuteFoldedDataDrivenTestsAsync.
+        // Any change made here should be reflected in TestMethodRunner.TryExecuteFoldedDataDrivenTestsAsync as well.
+        IEnumerable<object?[]> dataEnumerable = dataSource.GetData(methodInfo);
         string? testDataSourceIgnoreMessage = (dataSource as ITestDataSourceIgnoreCapability)?.IgnoreMessage;
 
-        if (!data.Any())
-        {
-            if (!MSTestSettings.CurrentSettings.ConsiderEmptyDataSourceAsInconclusive)
-            {
-                throw dataSource.GetExceptionForEmptyDataSource(methodInfo);
-            }
-
-            UnitTestElement discoveredTest = test.Clone();
-            // Make the test not data driven, because it had no data.
-            discoveredTest.TestMethod.DataType = DynamicDataType.None;
-            discoveredTest.TestMethod.TestDataSourceIgnoreMessage = testDataSourceIgnoreMessage;
-            discoveredTest.TestMethod.DisplayName = dataSource.GetDisplayName(methodInfo, null) ?? discoveredTest.TestMethod.DisplayName;
-            tests.Add(discoveredTest);
-
-            return true;
-        }
-
         var discoveredTests = new List<UnitTestElement>();
+        bool dataSourceHasData = false;
 
-        foreach (object?[] dataOrTestDataRow in data)
+        foreach (object?[] dataOrTestDataRow in dataEnumerable)
         {
+            dataSourceHasData = true;
             object?[] d = dataOrTestDataRow;
             ParameterInfo[] parameters = methodInfo.GetParameters();
             if (TestDataSourceHelpers.TryHandleITestDataRow(d, parameters, out d, out string? ignoreMessageFromTestDataRow, out string? displayNameFromTestDataRow, out IList<string>? testCategoriesFromTestDataRow))
@@ -365,6 +349,23 @@ private static bool TryUnfoldITestDataSource(ITestDataSource dataSource, UnitTes
             globalTestCaseIndex++;
         }
 
+        if (!dataSourceHasData)
+        {
+            if (!MSTestSettings.CurrentSettings.ConsiderEmptyDataSourceAsInconclusive)
+            {
+                throw dataSource.GetExceptionForEmptyDataSource(methodInfo);
+            }
+
+            UnitTestElement discoveredTest = test.Clone();
+            // Make the test not data driven, because it had no data.
+            discoveredTest.TestMethod.DataType = DynamicDataType.None;
+            discoveredTest.TestMethod.TestDataSourceIgnoreMessage = testDataSourceIgnoreMessage;
+            discoveredTest.TestMethod.DisplayName = dataSource.GetDisplayName(methodInfo, null) ?? discoveredTest.TestMethod.DisplayName;
+            tests.Add(discoveredTest);
+
+            return true;
+        }
+
         tests.AddRange(discoveredTests);
 
         return true;

src/Adapter/MSTestAdapter.PlatformServices/Execution/TestMethodRunner.cs

@@ -158,7 +158,7 @@ internal async Task<TestResult[]> RunTestMethodAsync()
         // In case of data driven, set parent info in results.
         if (isDataDriven)
         {
-            results = UpdateResultsWithParentInfo(results);
+            UpdateResultsWithParentInfo(results);
         }
 
         // Set a result in case no result is present.
@@ -205,13 +205,25 @@ private async Task<bool> TryExecuteFoldedDataDrivenTestsAsync(List<TestResult> r
                 continue;
             }
 
-            IEnumerable<object?[]>? dataSource;
+            // This code is to execute tests. To discover the tests code is in AssemblyEnumerator.TryUnfoldITestDataSource.
+            // Any change made here should be reflected in AssemblyEnumerator.TryUnfoldITestDataSource as well.
+            bool dataSourceHasData = false;
+            foreach (object?[] data in testDataSource.GetData(_testMethodInfo.MethodInfo))
+            {
+                dataSourceHasData = true;
+                try
+                {
+                    TestResult[] testResults = await ExecuteTestWithDataSourceAsync(testDataSource, data, actualDataAlreadyHandledDuringDiscovery: false).ConfigureAwait(false);
 
-            // This code is to execute tests. To discover the tests code is in AssemblyEnumerator.ProcessTestDataSourceTests.
-            // Any change made here should be reflected in AssemblyEnumerator.ProcessTestDataSourceTests as well.
-            dataSource = testDataSource.GetData(_testMethodInfo.MethodInfo);
+                    results.AddRange(testResults);
+                }
+                finally
+                {
+                    _testMethodInfo.SetArguments(null);
+                }
+            }
 
-            if (!dataSource.Any())
+            if (!dataSourceHasData)
             {
                 if (!MSTestSettings.CurrentSettings.ConsiderEmptyDataSourceAsInconclusive)
                 {
@@ -223,21 +235,6 @@ private async Task<bool> TryExecuteFoldedDataDrivenTestsAsync(List<TestResult> r
                     Outcome = UnitTestOutcome.Inconclusive,
                 };
                 results.Add(inconclusiveResult);
-                continue;
-            }
-
-            foreach (object?[] data in dataSource)
-            {
-                try
-                {
-                    TestResult[] testResults = await ExecuteTestWithDataSourceAsync(testDataSource, data, actualDataAlreadyHandledDuringDiscovery: false).ConfigureAwait(false);
-
-                    results.AddRange(testResults);
-                }
-                finally
-                {
-                    _testMethodInfo.SetArguments(null);
-                }
             }
         }
 
@@ -246,8 +243,7 @@ private async Task<bool> TryExecuteFoldedDataDrivenTestsAsync(List<TestResult> r
 
     private async Task ExecuteTestFromDataSourceAttributeAsync(List<TestResult> results)
     {
-        Stopwatch watch = new();
-        watch.Start();
+        var watch = Stopwatch.StartNew();
 
         try
         {
@@ -440,39 +436,24 @@ private static UnitTestOutcome GetAggregateOutcome(List<TestResult> results)
 
         // Get aggregate outcome.
         UnitTestOutcome aggregateOutcome = results[0].Outcome;
-        foreach (TestResult result in results)
+        for (int i = 1; i < results.Count; i++)
         {
-            aggregateOutcome = aggregateOutcome.GetMoreImportantOutcome(result.Outcome);
+            aggregateOutcome = aggregateOutcome.GetMoreImportantOutcome(results[i].Outcome);
         }
 
         return aggregateOutcome;
     }
 
     /// <summary>
-    /// Updates given results with parent info if results are greater than 1.
-    /// Add parent results as first result in updated result.
+    /// Updates each given result with new execution and parent execution identifiers.
     /// </summary>
     /// <param name="results">Results.</param>
-    /// <returns>Updated results which contains parent result as first result. All other results contains parent result info.</returns>
-    private static List<TestResult> UpdateResultsWithParentInfo(List<TestResult> results)
+    private static void UpdateResultsWithParentInfo(List<TestResult> results)
     {
-        // Return results in case there are no results.
-        if (results.Count == 0)
-        {
-            return results;
-        }
-
-        // UpdatedResults contain parent result at first position and remaining results has parent info updated.
-        var updatedResults = new List<TestResult>();
-
         foreach (TestResult result in results)
         {
             result.ExecutionId = Guid.NewGuid();
             result.ParentExecId = Guid.NewGuid();
-
-            updatedResults.Add(result);
         }
-
-        return updatedResults;
     }
 }

src/Adapter/MSTestAdapter.PlatformServices/Execution/UnitTestRunner.cs

@@ -269,12 +269,12 @@ private static async Task<TestResult> RunAssemblyInitializeIfNeededAsync(TestMet
 
     private static async Task<TestResult?> RunAssemblyCleanupAsync(ITestContext testContext, TypeCache typeCache, TestResult[] results)
     {
+        var testContextImpl = testContext as TestContextImplementation;
         IEnumerable<TestAssemblyInfo> assemblyInfoCache = typeCache.AssemblyInfoListWithExecutableCleanupMethods;
         foreach (TestAssemblyInfo assemblyInfo in assemblyInfoCache)
         {
             TestFailedException? ex = await assemblyInfo.ExecuteAssemblyCleanupAsync(testContext.Context).ConfigureAwait(false);
 
-            var testContextImpl = testContext as TestContextImplementation;
             if (ex is not null)
             {
                 return new TestResult()
@@ -314,32 +314,25 @@ private static bool IsTestMethodRunnable(
         [NotNullWhen(false)] out TestResult[]? notRunnableResult)
     {
         // If the specified TestMethod could not be found, return a NotFound result.
-        if (testMethodInfo == null)
+        if (testMethodInfo is null)
         {
-            {
-                notRunnableResult =
-                [
-                    new TestResult
-                    {
-                        Outcome = UnitTestOutcome.NotFound,
-                        IgnoreReason = string.Format(CultureInfo.CurrentCulture, Resource.TestNotFound, testMethod.Name),
-                    },
-                ];
-                return false;
-            }
+            notRunnableResult =
+            [
+                new TestResult
+                {
+                    Outcome = UnitTestOutcome.NotFound,
+                    IgnoreReason = string.Format(CultureInfo.CurrentCulture, Resource.TestNotFound, testMethod.Name),
+                },
+            ];
+            return false;
         }
 
         bool shouldIgnoreClass = testMethodInfo.Parent.ClassType.IsIgnored(out string? ignoreMessageOnClass);
         bool shouldIgnoreMethod = testMethodInfo.MethodInfo.IsIgnored(out string? ignoreMessageOnMethod);
 
-        string? ignoreMessage = ignoreMessageOnClass;
-        if (StringEx.IsNullOrEmpty(ignoreMessage) && shouldIgnoreMethod)
-        {
-            ignoreMessage = ignoreMessageOnMethod;
-        }
-
         if (shouldIgnoreClass || shouldIgnoreMethod)
         {
+            string? ignoreMessage = shouldIgnoreMethod && StringEx.IsNullOrEmpty(ignoreMessageOnClass) ? ignoreMessageOnMethod : ignoreMessageOnClass;
             notRunnableResult =
                 [TestResult.CreateIgnoredResult(ignoreMessage)];
             return false;