diff --git a/src/Tests/Microsoft.Diagnostics.Monitoring.Tool.UnitTests/CollectionRulePipelineTests.cs b/src/Tests/Microsoft.Diagnostics.Monitoring.Tool.UnitTests/CollectionRulePipelineTests.cs
index 4ba492b9b75..80ae91ecfd5 100644
--- a/src/Tests/Microsoft.Diagnostics.Monitoring.Tool.UnitTests/CollectionRulePipelineTests.cs
+++ b/src/Tests/Microsoft.Diagnostics.Monitoring.Tool.UnitTests/CollectionRulePipelineTests.cs
@@ -62,7 +62,7 @@ public Task CollectionRulePipeline_StartupTriggerTest(TargetFrameworkMoniker app
// Register first callback before pipeline starts. This callback should be completed before
// the pipeline finishes starting.
- Task callback1Task = callbackService.WaitWithCancellationAsync(cancellationSource.Token);
+ Task callback1Task = await callbackService.StartWaitForCallbackAsync(cancellationSource.Token);
// Startup trigger will cause the the pipeline to complete the start phase
// after the action list has been completed.
@@ -72,12 +72,11 @@ public Task CollectionRulePipeline_StartupTriggerTest(TargetFrameworkMoniker app
// completed because it was registered after the pipeline had finished starting. Since
// the action list is only ever executed once and is executed before the pipeline finishes
// starting, thus subsequent invocations of the action list should not occur.
- Task callback2Task = callbackService.WaitWithCancellationAsync(cancellationSource.Token);
+ Task callback2Task = await callbackService.StartWaitForCallbackAsync(cancellationSource.Token);
// Since the action list was completed before the pipeline finished starting,
// the action should have invoked it's callback.
- using CancellationTokenSource callbackCancellationSource = new(TimeSpan.FromMilliseconds(50));
- await callback1Task.WithCancellation(callbackCancellationSource.Token);
+ await callback1Task.WithCancellation(cancellationSource.Token);
// Regardless of the action list constraints, the pipeline should have only
// executed the action list once due to the use of a startup trigger.
@@ -132,7 +131,7 @@ public Task CollectionRulePipeline_EventCounterTriggerTest(TargetFrameworkMonike
// Register first callback before pipeline starts. This callback should be completed after
// the pipeline finishes starting.
- Task callbackTask = callbackService.WaitWithCancellationAsync(cancellationSource.Token);
+ Task callbackTask = await callbackService.StartWaitForCallbackAsync(cancellationSource.Token);
// Start pipeline with EventCounter trigger.
await pipeline.StartAsync(cancellationSource.Token);
@@ -140,7 +139,7 @@ public Task CollectionRulePipeline_EventCounterTriggerTest(TargetFrameworkMonike
await runner.SendCommandAsync(TestAppScenarios.SpinWait.Commands.StartSpin);
// This should not complete until the trigger conditions are satisfied for the first time.
- await callbackTask;
+ await callbackTask.WithCancellation(cancellationSource.Token);
VerifyExecutionCount(callbackService, 1);
diff --git a/src/Tests/Microsoft.Diagnostics.Monitoring.Tool.UnitTests/CollectionRules/Actions/CallbackAction.cs b/src/Tests/Microsoft.Diagnostics.Monitoring.Tool.UnitTests/CollectionRules/Actions/CallbackAction.cs
index 6901c488497..250073b566d 100644
--- a/src/Tests/Microsoft.Diagnostics.Monitoring.Tool.UnitTests/CollectionRules/Actions/CallbackAction.cs
+++ b/src/Tests/Microsoft.Diagnostics.Monitoring.Tool.UnitTests/CollectionRules/Actions/CallbackAction.cs
@@ -70,7 +70,22 @@ public async Task NotifyListeners(CancellationToken token)
}
}
- public async Task WaitWithCancellationAsync(CancellationToken token)
+ ///
+ /// Registers a callback with the Callback action that will complete when
+ /// the Callback action is invoked.
+ ///
+ ///
+ /// A that completes when the callback has finished
+ /// being registered. The inner will complete when the callback
+ /// is invoked.
+ ///
+ ///
+ /// Await this method to wait for the callback to be registered; await the inner
+ /// task to wait for the callback to be invoked. The parameter
+ /// only cancels registration if the registration has not completed; it does not cancel
+ /// the inner task that represents the callback invocation.
+ ///
+ public async Task StartWaitForCallbackAsync(CancellationToken token)
{
int id = _nextId++;
string name = $"Callback{id}";
@@ -89,7 +104,7 @@ public async Task WaitWithCancellationAsync(CancellationToken token)
_entriesSemaphore.Release();
}
- await entry.WithCancellation(token);
+ return entry.Task;
}
public IReadOnlyCollection ExecutionTimestamps
@@ -126,12 +141,7 @@ public void Complete()
_outputHelper.WriteLine("[Callback] End completing {0}.", _name);
}
- public async Task WithCancellation(CancellationToken token)
- {
- _outputHelper.WriteLine("[Test] Begin waiting for {0} completion.", _name);
- await _completionSource.WithCancellation(token);
- _outputHelper.WriteLine("[Test] End waiting for {0} completion.", _name);
- }
+ public Task Task => _completionSource.Task;
}
}
}