Skip to content

Commit ac903bb

Browse files
CopilotEvangelink
Copilot
and
Evangelink
committed
Add timeout parameter to TaskRunTestMethodAttribute
- Added timeout parameter to TaskRunTestMethodAttribute constructors - Implemented timeout handling with Task.WhenAny for non-cooperative timeout - Added STA thread handling to respect thread apartment state - Updated code fix to extract timeout from Timeout attribute and pass to TaskRunTestMethod - Updated tests to use TaskRunTestMethod(5000) instead of separate Timeout attribute - Updated PublicAPI.Unshipped.txt with new signatures Addresses comments from @Youssef1313 and @Evangelink about needing timeout parameter and STA thread handling. Co-authored-by: Evangelink <11340282+Evangelink@users.noreply.github.com>
1 parent cb78ce5 commit ac903bb

File tree

4 files changed

+153
-45
lines changed

4 files changed

+153
-45
lines changed

‎src/Analyzers/MSTest.Analyzers.CodeFixes/UseCooperativeCancellationForTimeoutFixer.cs‎

Lines changed: 39 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -134,40 +134,68 @@ private static async Task<Document> AddCooperativeCancellationAsync(Document doc
134134
private static async Task<Document> ReplaceWithTaskRunTestMethodAsync(Document document, MethodDeclarationSyntax methodDeclaration, CancellationToken cancellationToken)
135135
{
136136
DocumentEditor editor = await DocumentEditor.CreateAsync(document, cancellationToken).ConfigureAwait(false);
137+
SyntaxNode root = await document.GetSyntaxRootAsync(cancellationToken).ConfigureAwait(false);
137138

138-
// Find the TestMethod attribute
139+
if (root is null)
140+
{
141+
return document;
142+
}
143+
144+
// Find the TestMethod and Timeout attributes
139145
AttributeSyntax? testMethodAttribute = null;
146+
AttributeSyntax? timeoutAttribute = null;
147+
int timeoutValue = 0;
148+
140149
foreach (AttributeListSyntax attributeList in methodDeclaration.AttributeLists)
141150
{
142151
foreach (AttributeSyntax attribute in attributeList.Attributes)
143152
{
144153
string attributeName = attribute.Name.ToString();
154+
145155
if (attributeName is "TestMethod" or "TestMethodAttribute")
146156
{
147157
testMethodAttribute = attribute;
148-
break;
149158
}
150-
}
151-
152-
if (testMethodAttribute is not null)
153-
{
154-
break;
159+
else if (attributeName is "Timeout" or "TimeoutAttribute")
160+
{
161+
timeoutAttribute = attribute;
162+
163+
// Extract timeout value from the first argument
164+
if (attribute.ArgumentList?.Arguments.Count > 0)
165+
{
166+
var firstArg = attribute.ArgumentList.Arguments[0];
167+
if (firstArg.Expression is LiteralExpressionSyntax literalExpr &&
168+
literalExpr.Token.Value is int value)
169+
{
170+
timeoutValue = value;
171+
}
172+
}
173+
}
155174
}
156175
}
157176

158-
if (testMethodAttribute is null)
177+
if (testMethodAttribute is null || timeoutAttribute is null)
159178
{
160-
// No TestMethod attribute found, return unchanged document
179+
// Can't apply the fix without both attributes
161180
return document;
162181
}
163182

164-
// Create the new TaskRunTestMethod attribute preserving any arguments
183+
// Create the new TaskRunTestMethod attribute with timeout parameter
184+
AttributeArgumentSyntax timeoutArg = SyntaxFactory.AttributeArgument(
185+
SyntaxFactory.LiteralExpression(
186+
SyntaxKind.NumericLiteralExpression,
187+
SyntaxFactory.Literal(timeoutValue)));
188+
165189
AttributeSyntax newAttribute = SyntaxFactory.Attribute(
166190
SyntaxFactory.IdentifierName("TaskRunTestMethod"),
167-
testMethodAttribute.ArgumentList);
191+
SyntaxFactory.AttributeArgumentList(
192+
SyntaxFactory.SingletonSeparatedList(timeoutArg)));
168193

169194
// Replace the TestMethod attribute with TaskRunTestMethod
170195
editor.ReplaceNode(testMethodAttribute, newAttribute);
196+
197+
// Remove the Timeout attribute
198+
editor.RemoveNode(timeoutAttribute);
171199

172200
return editor.GetChangedDocument();
173201
}

‎src/TestFramework/TestFramework/Attributes/TestMethod/TaskRunTestMethodAttribute.cs‎

Lines changed: 105 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -4,74 +4,151 @@
44
namespace Microsoft.VisualStudio.TestTools.UnitTesting;
55

66
/// <summary>
7-
/// The TaskRun test method attribute.
7+
/// Test method attribute that runs tests in a Task.Run with non-cooperative timeout handling.
88
/// </summary>
99
/// <remarks>
1010
/// <para>
11-
/// This attribute is designed to handle test method execution with timeout by running the test code within a <see cref="Task.Run"/>.
12-
/// This allows the test runner to stop watching the task in case of timeout, preventing dangling tasks that can lead to
13-
/// confusion or errors because the test method is still running in the background.
11+
/// This attribute runs test methods in <see cref="Task.Run"/> and implements non-cooperative timeout handling.
12+
/// When a timeout occurs, the test is marked as timed out and the test runner stops awaiting the task,
13+
/// allowing it to complete in the background. This prevents blocking but may lead to dangling tasks.
1414
/// </para>
1515
/// <para>
16-
/// When a timeout occurs:
17-
/// <list type="bullet">
18-
/// <item><description>The test is marked as timed out.</description></item>
19-
/// <item><description>The cancellation token from <see cref="TestContext.CancellationToken"/> is canceled.</description></item>
20-
/// <item><description>The test runner stops awaiting the test task, allowing it to complete in the background.</description></item>
21-
/// </list>
22-
/// </para>
23-
/// <para>
24-
/// For best results, test methods should observe the cancellation token and cancel cooperatively.
25-
/// If the test method does not handle cancellation properly, the task may continue running after the timeout,
26-
/// which can still lead to issues, but the test runner will not block waiting for it to complete.
16+
/// For cooperative timeout handling where tests are awaited until completion, use <see cref="TimeoutAttribute"/>
17+
/// with <c>CooperativeCancellation = true</c> instead.
2718
/// </para>
2819
/// </remarks>
20+
[AttributeUsage(AttributeTargets.Method, Inherited = false)]
2921
public sealed class TaskRunTestMethodAttribute : TestMethodAttribute
3022
{
3123
private readonly TestMethodAttribute? _testMethodAttribute;
3224

3325
/// <summary>
3426
/// Initializes a new instance of the <see cref="TaskRunTestMethodAttribute"/> class.
3527
/// </summary>
36-
public TaskRunTestMethodAttribute([CallerFilePath] string callerFilePath = "", [CallerLineNumber] int callerLineNumber = -1)
28+
/// <param name="timeout">The timeout in milliseconds. If not specified or 0, no timeout is applied.</param>
29+
public TaskRunTestMethodAttribute(int timeout = 0, [CallerFilePath] string callerFilePath = "", [CallerLineNumber] int callerLineNumber = -1)
3730
: base(callerFilePath, callerLineNumber)
3831
{
32+
Timeout = timeout;
3933
}
4034

4135
/// <summary>
4236
/// Initializes a new instance of the <see cref="TaskRunTestMethodAttribute"/> class.
4337
/// This constructor is intended to be called by test class attributes to wrap an existing test method attribute.
4438
/// </summary>
4539
/// <param name="testMethodAttribute">The wrapped test method.</param>
46-
public TaskRunTestMethodAttribute(TestMethodAttribute testMethodAttribute)
40+
/// <param name="timeout">The timeout in milliseconds. If not specified or 0, no timeout is applied.</param>
41+
public TaskRunTestMethodAttribute(TestMethodAttribute testMethodAttribute, int timeout = 0)
4742
: base(testMethodAttribute.DeclaringFilePath, testMethodAttribute.DeclaringLineNumber ?? -1)
48-
=> _testMethodAttribute = testMethodAttribute;
43+
{
44+
_testMethodAttribute = testMethodAttribute;
45+
Timeout = timeout;
46+
}
4947

5048
/// <summary>
51-
/// Executes a test method by wrapping it in a <see cref="Task.Run"/> to allow timeout handling.
49+
/// Gets the timeout in milliseconds.
50+
/// </summary>
51+
public int Timeout { get; }
52+
53+
/// <summary>
54+
/// Executes a test method with non-cooperative timeout handling.
5255
/// </summary>
5356
/// <param name="testMethod">The test method to execute.</param>
5457
/// <returns>An array of TestResult objects that represent the outcome(s) of the test.</returns>
5558
public override async Task<TestResult[]> ExecuteAsync(ITestMethod testMethod)
5659
{
5760
if (_testMethodAttribute is not null)
5861
{
59-
return await ExecuteWithTaskRunAsync(() => _testMethodAttribute.ExecuteAsync(testMethod)).ConfigureAwait(false);
62+
return await ExecuteWithTimeoutAsync(() => _testMethodAttribute.ExecuteAsync(testMethod), testMethod).ConfigureAwait(false);
6063
}
6164

62-
return await ExecuteWithTaskRunAsync(async () =>
65+
return await ExecuteWithTimeoutAsync(async () =>
6366
{
6467
TestResult result = await testMethod.InvokeAsync(null).ConfigureAwait(false);
6568
return new[] { result };
66-
}).ConfigureAwait(false);
69+
}, testMethod).ConfigureAwait(false);
6770
}
6871

69-
private static async Task<TestResult[]> ExecuteWithTaskRunAsync(Func<Task<TestResult[]>> executeFunc)
72+
private async Task<TestResult[]> ExecuteWithTimeoutAsync(Func<Task<TestResult[]>> executeFunc, ITestMethod testMethod)
7073
{
71-
// Run the test method in Task.Run so that we can stop awaiting it on timeout
72-
// while allowing it to complete in the background
73-
Task<TestResult[]> testTask = Task.Run(executeFunc);
74-
TestResult[] results = await testTask.ConfigureAwait(false);
75-
return results;
74+
if (Timeout <= 0)
75+
{
76+
// No timeout, run directly with Task.Run
77+
return await RunOnThreadPoolOrCustomThreadAsync(executeFunc).ConfigureAwait(false);
78+
}
79+
80+
// Run with timeout
81+
Task<TestResult[]> testTask = RunOnThreadPoolOrCustomThreadAsync(executeFunc);
82+
Task completedTask = await Task.WhenAny(testTask, Task.Delay(Timeout)).ConfigureAwait(false);
83+
84+
if (completedTask == testTask)
85+
{
86+
// Test completed before timeout
87+
return await testTask.ConfigureAwait(false);
88+
}
89+
90+
// Timeout occurred - return timeout result and let task continue in background
91+
return
92+
[
93+
new TestResult
94+
{
95+
Outcome = UnitTestOutcome.Timeout,
96+
TestFailureException = new TestFailedException(
97+
UnitTestOutcome.Timeout,
98+
string.Format(
99+
CultureInfo.InvariantCulture,
100+
"Test '{0}.{1}' exceeded timeout of {2}ms.",
101+
testMethod.TestClassName,
102+
testMethod.TestMethodName,
103+
Timeout)),
104+
},
105+
];
106+
}
107+
108+
private static Task<TestResult[]> RunOnThreadPoolOrCustomThreadAsync(Func<Task<TestResult[]>> executeFunc)
109+
{
110+
// Check if we need to handle STA threading
111+
// If current thread is STA and we're on Windows, create a new STA thread
112+
// Otherwise, use Task.Run (thread pool)
113+
#if NETFRAMEWORK
114+
if (Thread.CurrentThread.GetApartmentState() == ApartmentState.STA)
115+
{
116+
return RunOnSTAThreadAsync(executeFunc);
117+
}
118+
#else
119+
if (RuntimeInformation.IsOSPlatform(OSPlatform.Windows) &&
120+
Thread.CurrentThread.GetApartmentState() == ApartmentState.STA)
121+
{
122+
return RunOnSTAThreadAsync(executeFunc);
123+
}
124+
#endif
125+
126+
// Use thread pool for non-STA scenarios
127+
return Task.Run(executeFunc);
128+
}
129+
130+
private static Task<TestResult[]> RunOnSTAThreadAsync(Func<Task<TestResult[]>> executeFunc)
131+
{
132+
var tcs = new TaskCompletionSource<TestResult[]>();
133+
var thread = new Thread(() =>
134+
{
135+
try
136+
{
137+
TestResult[] result = executeFunc().GetAwaiter().GetResult();
138+
tcs.SetResult(result);
139+
}
140+
catch (Exception ex)
141+
{
142+
tcs.SetException(ex);
143+
}
144+
})
145+
{
146+
Name = "TaskRunTestMethodAttribute STA thread",
147+
};
148+
149+
thread.SetApartmentState(ApartmentState.STA);
150+
thread.Start();
151+
152+
return tcs.Task;
76153
}
77154
}

‎src/TestFramework/TestFramework/PublicAPI/PublicAPI.Unshipped.txt‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
#nullable enable
22
Microsoft.VisualStudio.TestTools.UnitTesting.TaskRunTestMethodAttribute
33
Microsoft.VisualStudio.TestTools.UnitTesting.TaskRunTestMethodAttribute.ExecuteAsync(Microsoft.VisualStudio.TestTools.UnitTesting.ITestMethod testMethod) -> System.Threading.Tasks.Task<Microsoft.VisualStudio.TestTools.UnitTesting.TestResult[]!>!
4-
Microsoft.VisualStudio.TestTools.UnitTesting.TaskRunTestMethodAttribute.TaskRunTestMethodAttribute(Microsoft.VisualStudio.TestTools.UnitTesting.TestMethodAttribute! testMethodAttribute) -> void
5-
Microsoft.VisualStudio.TestTools.UnitTesting.TaskRunTestMethodAttribute.TaskRunTestMethodAttribute(string! callerFilePath = "", int callerLineNumber = -1) -> void
4+
Microsoft.VisualStudio.TestTools.UnitTesting.TaskRunTestMethodAttribute.TaskRunTestMethodAttribute(int timeout = 0, string! callerFilePath = "", int callerLineNumber = -1) -> void
5+
Microsoft.VisualStudio.TestTools.UnitTesting.TaskRunTestMethodAttribute.TaskRunTestMethodAttribute(Microsoft.VisualStudio.TestTools.UnitTesting.TestMethodAttribute! testMethodAttribute, int timeout = 0) -> void
6+
Microsoft.VisualStudio.TestTools.UnitTesting.TaskRunTestMethodAttribute.Timeout.get -> int

‎test/UnitTests/MSTest.Analyzers.UnitTests/UseCooperativeCancellationForTimeoutAnalyzerTests.cs‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -464,8 +464,7 @@ public void MyTestMethod()
464464
[TestClass]
465465
public class MyTestClass
466466
{
467-
[TaskRunTestMethod]
468-
[Timeout(5000)]
467+
[TaskRunTestMethod(5000)]
469468
public void MyTestMethod()
470469
{
471470
}
@@ -488,8 +487,7 @@ public async Task WhenTaskRunTestMethodAttributeWithTimeout_NoDiagnostic()
488487
[TestClass]
489488
public class MyTestClass
490489
{
491-
[TaskRunTestMethod]
492-
[Timeout(5000)]
490+
[TaskRunTestMethod(5000)]
493491
public void MyTestMethod()
494492
{
495493
}
@@ -498,4 +496,8 @@ public void MyTestMethod()
498496

499497
await VerifyCS.VerifyCodeFixAsync(code, code);
500498
}
499+
""";
500+
501+
await VerifyCS.VerifyCodeFixAsync(code, code);
502+
}
501503
}

0 commit comments

Comments
 (0)