How to correctly wait for callback with timeout fallback

Question

I am testing web socket subscriptions in my tests and I would like to wait for response from callback and then end the test, if no response is received after timeout end the test.

This is what I have now (simplified) but I am not sure if its the way how to do it.

public async Task WaitForPing()
{
  var cancellationTokenSource = new CancellationTokenSource(5_000);
  var pinged = false;
  using var _ = Client.OnPing(_ =>
  {
    pinged = true;
    cancellationTokenSource.Cancel();
  }

  await Client.Run();

  await Task
    .Delay(-1, cancellationTokenSource.Token)
    .ContinueWith(_ => { }, CancellationToken.None);
  Assert(pinged);
}

Answer 1

A proper method should be like that:

static Task<string> WaitForResponseAsync(CancellationTokenSource cancellationTokenSource = default)
{
    var tcs = new TaskCompletionSource<string>();

    // Register a method that throws an exception when task cancelled.
    cancellationTokenSource.Token.Register(()=> throw new Exception("Timed out!"));

    // Replace this task with your async operation. Like OnPing(_ => ...
    Task.Run(async () =>
    {
        await Task.Delay(30_000); // Response will be received after 30 seconds
        tcs.SetResult("Hello World");
    });

    return tcs.Task; // Return awaitable task
}

And place where you call that method:

try
{
    Console.WriteLine(await WaitForResponseAsync(new CancellationTokenSource(5_000))); // Time out is 5 seconds
}
catch (Exception ex)
{
    Console.WriteLine(ex.ToString());
}

---

Don't worry about the task in a task, you should replace it with your event. So if I customize it according to your method, it should be something like that:

Task WaitForPing(CancellationTokenSource cancellationTokenSource)
{
    var tcs = new TaskCompletionSource();

    cancellationTokenSource.Token.Register(() => throw new Exception("Timed out"));

    // Or just call SetResult to finish the task before completed without exception: 
    // cancellationTokenSource.Token.Register(() => tcs.SetResult());
    // (Personally I do not recommend this one)

    using var _ = Client.OnPing(_ =>
    {
        tcs.SetResult();
    };

    return tcs.Task;
}