Better handling of CTS and timeout tokens

src/Adapter/MSTest.TestAdapter/Execution/TestMethodInfo.cs

@@ -22,7 +22,7 @@ namespace Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter.Execution;
 [Obsolete(Constants.PublicTypeObsoleteMessage)]
 #endif
 #endif
-public class TestMethodInfo : ITestMethod
+public class TestMethodInfo : ITestMethod, IDisposable
 {
     /// <summary>
     /// Specifies the timeout when it is not set in a test case.
@@ -32,6 +32,7 @@ public class TestMethodInfo : ITestMethod
     private object? _classInstance;
     private bool _isTestContextSet;
     private bool _isTestCleanupInvoked;
+    private CancellationTokenSource? _timeoutTokenSource;
 
     internal TestMethodInfo(
         MethodInfo testMethod,
@@ -482,6 +483,15 @@ private void RunTestCleanupMethod(TestResult result)
         {
             try
             {
+                // Reset the cancellation token source to avoid cancellation of cleanup methods because of the init or test method cancellation.
+                TestMethodOptions.TestContext!.Context.CancellationTokenSource = new CancellationTokenSource();
+
+                // If we are running with a method timeout, we need to cancel the cleanup when the overall timeout expires. If it already expired, nothing to do.
+                if (_timeoutTokenSource is { IsCancellationRequested: false })
+                {
+                    _timeoutTokenSource?.Token.Register(TestMethodOptions.TestContext.Context.CancellationTokenSource.Cancel);
+                }
+
                 // Test cleanups are called in the order of discovery
                 // Current TestClass -> Parent -> Grandparent
                 testCleanupException = testCleanupMethod is not null
@@ -782,18 +792,27 @@ private bool SetTestContext(object classInstance, TestResult result)
             // In most cases, exception will be TargetInvocationException with real exception wrapped
             // in the InnerException; or user code throws an exception.
             // It also seems that in rare cases the ex can be null.
-            Exception actualException = ex.InnerException ?? ex;
-            string exceptionMessage = actualException.GetFormattedExceptionMessage();
-            StackTraceInformation? stackTraceInfo = actualException.GetStackTraceInformation();
+            Exception realException = ex.GetRealException();
 
-            string errorMessage = string.Format(
-                CultureInfo.CurrentCulture,
-                Resource.UTA_InstanceCreationError,
-                TestClassName,
-                exceptionMessage);
+            if (realException.IsOperationCanceledExceptionFromToken(TestMethodOptions.TestContext!.Context.CancellationTokenSource.Token))
+            {
+                result.Outcome = UTF.UnitTestOutcome.Timeout;
+                result.TestFailureException = new TestFailedException(ObjectModelUnitTestOutcome.Timeout, string.Format(CultureInfo.CurrentCulture, Resource.Execution_Test_Timeout, TestMethodName));
+            }
+            else
+            {
+                string exceptionMessage = realException.GetFormattedExceptionMessage();
+                StackTraceInformation? stackTraceInfo = realException.GetStackTraceInformation();
 
-            result.Outcome = UTF.UnitTestOutcome.Failed;
-            result.TestFailureException = new TestFailedException(ObjectModelUnitTestOutcome.Failed, errorMessage, stackTraceInfo);
+                string errorMessage = string.Format(
+                    CultureInfo.CurrentCulture,
+                    Resource.UTA_InstanceCreationError,
+                    TestClassName,
+                    exceptionMessage);
+
+                result.Outcome = UTF.UnitTestOutcome.Failed;
+                result.TestFailureException = new TestFailedException(ObjectModelUnitTestOutcome.Failed, errorMessage, stackTraceInfo);
+            }
         }
 
         return classInstance;
@@ -811,12 +830,11 @@ private TestResult ExecuteInternalWithTimeout(object?[]? arguments)
 
         if (TestMethodOptions.TimeoutInfo.CooperativeCancellation)
         {
-            CancellationTokenSource? timeoutTokenSource = null;
             try
             {
-                timeoutTokenSource = new(TestMethodOptions.TimeoutInfo.Timeout);
-                timeoutTokenSource.Token.Register(TestMethodOptions.TestContext!.Context.CancellationTokenSource.Cancel);
-                if (timeoutTokenSource.Token.IsCancellationRequested)
+                _timeoutTokenSource = new(TestMethodOptions.TimeoutInfo.Timeout);
+                _timeoutTokenSource.Token.Register(TestMethodOptions.TestContext!.Context.CancellationTokenSource.Cancel);
+                if (_timeoutTokenSource.Token.IsCancellationRequested)
                 {
                     return new()
                     {
@@ -840,15 +858,16 @@ private TestResult ExecuteInternalWithTimeout(object?[]? arguments)
                         Outcome = UTF.UnitTestOutcome.Timeout,
                         TestFailureException = new TestFailedException(
                             ObjectModelUnitTestOutcome.Timeout,
-                            timeoutTokenSource.Token.IsCancellationRequested
+                            _timeoutTokenSource.Token.IsCancellationRequested
                                 ? string.Format(CultureInfo.CurrentCulture, Resource.Execution_Test_Timeout, TestMethodName)
                                 : string.Format(CultureInfo.CurrentCulture, Resource.Execution_Test_Cancelled, TestMethodName)),
                     };
                 }
             }
             finally
             {
-                timeoutTokenSource?.Dispose();
+                _timeoutTokenSource?.Dispose();
+                _timeoutTokenSource = null;
             }
         }
 
@@ -867,8 +886,7 @@ void ExecuteAsyncAction()
             }
         }
 
-        CancellationToken cancelToken = TestMethodOptions.TestContext!.Context.CancellationTokenSource.Token;
-        if (PlatformServiceProvider.Instance.ThreadOperations.Execute(ExecuteAsyncAction, TestMethodOptions.TimeoutInfo.Timeout, cancelToken))
+        if (PlatformServiceProvider.Instance.ThreadOperations.Execute(ExecuteAsyncAction, TestMethodOptions.TimeoutInfo.Timeout, TestMethodOptions.TestContext!.Context.CancellationTokenSource.Token))
         {
             if (failure != null)
             {
@@ -902,4 +920,7 @@ void ExecuteAsyncAction()
         RunTestCleanupMethod(timeoutResult);
         return timeoutResult;
     }
+
+    void IDisposable.Dispose()
+        => _timeoutTokenSource?.Dispose();
 }

src/Adapter/MSTest.TestAdapter/Extensions/ExceptionExtensions.cs

@@ -32,6 +32,10 @@ internal static Exception GetRealException(this Exception exception)
         return exception;
     }
 
+    internal static bool IsOperationCanceledExceptionFromToken(this Exception ex, CancellationToken cancellationToken)
+        => (ex is OperationCanceledException oce && oce.CancellationToken == cancellationToken)
+        || (ex is AggregateException aggregateEx && aggregateEx.InnerExceptions.OfType<OperationCanceledException>().Any(oce => oce.CancellationToken == cancellationToken));
+
     /// <summary>
     /// Get the exception message if available, empty otherwise.
     /// </summary>

src/Adapter/MSTest.TestAdapter/Helpers/FixtureMethodRunner.cs

@@ -16,6 +16,13 @@ internal static class FixtureMethodRunner
         Action action, CancellationTokenSource cancellationTokenSource, TimeoutInfo? timeoutInfo, MethodInfo methodInfo,
         IExecutionContextScope executionContextScope, string methodCanceledMessageFormat, string methodTimedOutMessageFormat)
     {
+        if (cancellationTokenSource.Token.IsCancellationRequested)
+        {
+            return new(
+                UnitTestOutcome.Timeout,
+                string.Format(CultureInfo.InvariantCulture, methodCanceledMessageFormat, methodInfo.DeclaringType!.FullName, methodInfo.Name));
+        }
+
         // If no timeout is specified, we can run the action directly. This avoids any overhead of creating a task/thread and
         // ensures that the execution context is preserved (as we run the action on the current thread).
         if (timeoutInfo is null)
@@ -27,9 +34,7 @@ internal static class FixtureMethodRunner
             }
             catch (Exception ex)
             {
-                Exception realException = ex.GetRealException();
-
-                if (realException is OperationCanceledException oce && oce.CancellationToken == cancellationTokenSource.Token)
+                if (ex.GetRealException().IsOperationCanceledExceptionFromToken(cancellationTokenSource.Token))
                 {
                     return new(
                         UnitTestOutcome.Timeout,
@@ -75,7 +80,7 @@ internal static class FixtureMethodRunner
                 action();
                 return null;
             }
-            catch (OperationCanceledException)
+            catch (Exception ex) when (ex.IsOperationCanceledExceptionFromToken(cancellationTokenSource.Token))
             {
                 // Ideally we would like to check that the token of the exception matches cancellationTokenSource but TestContext
                 // instances are not well defined so we have to handle the exception entirely.
@@ -148,9 +153,7 @@ internal static class FixtureMethodRunner
                     methodInfo.Name,
                     timeout));
         }
-        catch (Exception ex) when
-            (ex is OperationCanceledException
-            || (ex is AggregateException aggregateEx && aggregateEx.InnerExceptions.OfType<TaskCanceledException>().Any()))
+        catch (Exception ex) when (ex.IsOperationCanceledExceptionFromToken(cancellationTokenSource.Token))
         {
             return new(
                 UnitTestOutcome.Timeout,
@@ -168,7 +171,7 @@ internal static class FixtureMethodRunner
         }
     }
 
-    [System.Runtime.Versioning.SupportedOSPlatform("windows")]
+    [SupportedOSPlatform("windows")]
     private static TestFailedException? RunWithTimeoutAndCancellationWithSTAThread(
         Action action, CancellationTokenSource cancellationTokenSource, int timeout, MethodInfo methodInfo,
         IExecutionContextScope executionContextScope, string methodCanceledMessageFormat, string methodTimedOutMessageFormat)
@@ -212,12 +215,7 @@ internal static class FixtureMethodRunner
                     methodInfo.Name,
                     timeout));
         }
-        catch (Exception ex) when
-
-            // This exception occurs when the cancellation happens before the task is actually started.
-            ((ex is TaskCanceledException tce && tce.CancellationToken == cancellationTokenSource.Token)
-            || (ex is OperationCanceledException oce && oce.CancellationToken == cancellationTokenSource.Token)
-            || (ex is AggregateException aggregateEx && aggregateEx.InnerExceptions.OfType<TaskCanceledException>().Any()))
+        catch (Exception ex) when (ex.IsOperationCanceledExceptionFromToken(cancellationTokenSource.Token))
         {
             return new(
                 UnitTestOutcome.Timeout,

src/Analyzers/MSTest.Analyzers/TestContextShouldBeValidAnalyzer.cs

@@ -124,7 +124,7 @@ private static void CollectTestContextFieldsAssignedInConstructor(
         }
 
         if (operation is ISimpleAssignmentOperation assignmentOperation &&
-            assignmentOperation.Target is IMemberReferenceOperation { Member: IFieldSymbol { } candidateField } targetMemberReference &&
+            assignmentOperation.Target is IMemberReferenceOperation { Member: IFieldSymbol { } candidateField } &&
             assignmentOperation.Value is IParameterReferenceOperation parameterReference &&
             SymbolEqualityComparer.Default.Equals(parameterReference.Parameter, testContextParameter))
         {

src/TestFramework/TestFramework.Extensions/TestContext.cs

@@ -48,7 +48,7 @@ public abstract class TestContext
     /// <summary>
     /// Gets or sets the cancellation token source. This token source is canceled when test times out. Also when explicitly canceled the test will be aborted.
     /// </summary>
-    public virtual CancellationTokenSource CancellationTokenSource { get; protected set; } = new();
+    public virtual CancellationTokenSource CancellationTokenSource { get; protected internal set; } = new();
 
     public object?[]? TestData { get; protected set; }
 

src/TestFramework/TestFramework.Extensions/TestFramework.Extensions.csproj

@@ -81,6 +81,10 @@
     <Using Include="Polyfills" />
   </ItemGroup>
 
+  <ItemGroup>
+    <InternalsVisibleTo Include="Microsoft.VisualStudio.TestPlatform.MSTest.TestAdapter" Key="$(VsPublicKey)" />
+  </ItemGroup>
+
   <ItemGroup>
     <None Update="build\winui\MSTest.TestFramework.targets">
       <CopyToOutputDirectory>PreserveNewest</CopyToOutputDirectory>