Wait until all Task finish in unit test

Question

I have this class I want to unit test:

public class SomeClass
{
    public void Foo()
    {
        Bar();
    }

    private void Bar()
    {
        Task.Factory.StartNew(() =>
        {
            // Do something that takes some time (e.g. an HTTP request)
        });
    }
}

And this is how my unit test looks like:

[TestMethod]
public void TestFoo()
{
    // Arrange
    var obj = new SomeClass();

    // Act
    obj.Foo();
    obj.Foo();
    obj.Foo();

    // Assert
    /* I need something to wait on all tasks to finish */
    Assert.IsTrue(...);
}

So, I need to make the unit test thread wait until all tasks started in the Bar method have finished their job before starting my assertions.

IMPORTANT: I cannot change SomeClass

How can I do that?

Answer 1

One way to solve this problem is to define your own task scheduler in such a way that would allow you to keep track of the completion of your nested tasks. For example, you could define a scheduler that executes tasks synchronously, as below:

class SynchronousTaskScheduler : TaskScheduler
{
    protected override void QueueTask(Task task)
    {
        this.TryExecuteTask(task);
    }

    protected override bool TryExecuteTaskInline(Task task, bool wasPreviouslyQueued)
    {
        return this.TryExecuteTask(task);
    }

    protected override IEnumerable<Task> GetScheduledTasks()
    {
        yield break;
    }
}

Subsequently, create an instance of this synchronous task scheduler, and use it to execute a root task which, in turn, spawns all of your "hidden" tasks. Since nested tasks inherit the current task scheduler from their parent, all your inner tasks will also get run on our synchronous scheduler, implying that our outermost StartNew call will only return when all tasks complete.

TaskScheduler scheduler = new SynchronousTaskScheduler();

Task.Factory.StartNew(() =>
{
    // Arrange
    var obj = new SomeClass();

    // Act
    obj.Foo();
    obj.Foo();
    obj.Foo();
}, 
    CancellationToken.None,
    TaskCreationOptions.None,
    scheduler);

// Assert
/* I need something to wait on all tasks to finish */
Assert.IsTrue(...);

A downside to this approach is that you will lose all concurrency from your tasks; however, you could fix this by enhancing the custom scheduler to one which is concurrent but still allows you to track executing tasks.

Answer 2

Task.WaitAll(the, list, of, task, objects, you, need, to, wait, on);

If it's a void async method then you can't do it. The design is broken. They're only for fire and forget.

Answer 3

Not sure if you are allowed to make this change but I got it to work doing this:

namespace ParallelProgramming.Playground
{
    public class SomeClass
    {
        public Task Foo()
        {
            return Bar();
        }

        private static Task Bar()
        {
            return Task.Factory.StartNew(() =>
                {
                    Console.WriteLine("I fired off. Thread ID: {0}", Thread.CurrentThread.ManagedThreadId);
                    Thread.Sleep(5000);
                    return true; //or whatever else you want.
                });
        }
    }

    [TestClass]
    public class StackOverflow
    {
        [TestMethod]
        public void TestFoo()
        {
            // Arrange
            var obj = new SomeClass();

            var results = new ConcurrentBag<Task>(); 
            var waitForMe = Task.Factory.StartNew(() =>
                {
                    // Act
                    results.Add(obj.Foo());
                    results.Add(obj.Foo());
                    results.Add(obj.Foo());

                    return true;
                });


            Task.WaitAll(waitForMe);

            // Assert
            /* I need something to wait on all tasks to finish */
            Assert.IsTrue(waitForMe.Result);
            Assert.AreEqual(3, results.Count);
        }
    }
}