TPL in a Unit Test? GENIUS!

By | May 14, 2014

How often have you created a unit test that you knew was long running (test all ___ overnight), only to find that eventually you hit a time when it never stopped, and horked up the rest of the test run?

It’s happened to the best of us, I’m sure.

However, thanks to .Net’s beautiful TPL, this becomes insanely easy. Let’s check it out.

The first thing we need to realize is what we want to do. What that is, in this case, is the ability to cancel a test mid-run after we know it’s gone too long. It’s somewhat along the lines of deadlock prevention, but not quite as academic ;)

So how does TPL enable task cancellation? CancellationTokenSource. You start a task, pass it a token, then you can use that token to cancel the task mid-run. With that in mind, let’s just get down to brass tax.

   1: [TestMethod]

   2: [ExpectedException(typeof(OperationCanceledException))]

   3: public void TestMethod1()

   4: {

   5:     var cts = new CancellationTokenSource();

   6:     cts.Token.Register(() =>

   7:     {

   8:         Console.WriteLine("stuff to print to console if the test runs too long");

   9:     });

  10:     Task.Factory.StartNew(() =>

  11:     {

  12:         cts.CancelAfter(4000);

  13:         Thread.Sleep(5000);

  14:     })

  15:     .ContinueWith(t =>

  16:     {

  17:         if (t.Exception != null)

  18:         {

  19:             throw t.Exception;

  20:         }

  21:  

  22:         Console.WriteLine("Test completed w/in the alotted time");

  23:     })

  24:     .Wait(cts.Token);

  25: }

Let’s review the code.

First we create the CTS object.
Next, register some work to be done if/when the task ends up getting cancelled.
After that, fire off the actual work w/in a Task object.
Tell the CTS how long to wait before cancelling.
Do the long-running work.
After the task is done, continue by analyzing and re-throwing any exception aggregated by the task, or doing any additional work if you so choose.
Finally, wait on all this work based on the token. Either the task completes fully, or the token’s cancellation is tripped.

When it’s all said and done, the above UT passes (because it encounters an exception) as well as the following UT, because it finishes w/in the alotted time:

   1: [TestMethod]

   2: public void TestMethod2()

   3: {

   4:     var cts = new CancellationTokenSource();

   5:     cts.Token.Register(() =>

   6:     {

   7:         Console.WriteLine("stuff to print to console if the test runs too long");

   8:     });

   9:     Task.Factory.StartNew(() =>

  10:     {

  11:         cts.CancelAfter(6000);

  12:         Thread.Sleep(5000);

  13:     })

  14:     .ContinueWith(t =>

  15:     {

  16:         if (t.Exception != null)

  17:         {

  18:             throw t.Exception;

  19:         }

  20:  

  21:         Console.WriteLine("Test completed w/in the alotted time");

  22:     })

  23:     .Wait(cts.Token);

  24: }

Here’s the Test Run output for TestMethod1:

Test Name:    TestMethod1
Test Outcome:    Passed
Result StandardOutput:    stuff to print to console if the test runs too long

And for TestMethod2:

Test Name:    TestMethod2
Test Outcome:    Passed
Result StandardOutput:    Test completed w/in the alotted time

Now, I’m sure there are many different ways to solve this problem w/ TPL and Cancellation Tokens, and I’m certainly all ears. This one certainly looks like it’s going to work well for me in my situations. Let yours rip in the comments!