Improve StartWithContinuations test in AsyncType tests

tests/FSharp.Core.UnitTests/FSharp.Core/Microsoft.FSharp.Control/AsyncType.fs

@@ -11,74 +11,48 @@ open Xunit
 open System.Threading
 open System.Threading.Tasks
 
-type RunWithContinuationsTest_WhatToDo =
-    | Exit
-    | Cancel
-    | Throw
+module AsyncType =
 
-type AsyncType() =
+    type ExpectedContinuation = Success | Exception | Cancellation
 
-    let ignoreSynchCtx f =
-        f ()
+    [<Fact>]
+    let startWithContinuations() =
 
-    [<VolatileField>]
-    let mutable spinloop = true
-
-    let waitASec (t:Task) =
-        let result = t.Wait(TimeSpan(hours=0,minutes=0,seconds=1))
-        Assert.True(result, "Task did not finish after waiting for a second.")
+        let cont actual expected _ =
+            if expected <> actual then
+                failwith $"expected {expected} continuation, but ran {actual}"
 
-    [<Fact>]
-    member _.StartWithContinuations() =
+        let onSuccess = cont Success
+        let onExeption = cont Exception
+        let onCancellation = cont Cancellation
 
-        let mutable whatToDo = Exit
+        let expect expected cancellationToken computation =
+            Async.StartWithContinuations(computation, onSuccess expected, onExeption expected, onCancellation expected, ?cancellationToken = cancellationToken)
 
-        let asyncWorkflow() =
-            async {
-                let currentState = whatToDo
-
-                // Act
-                let result =
-                    match currentState with
-                    | Exit   -> 1
-                    | Cancel -> Async.CancelDefaultToken()
-                                sleep(1 * 1000)
-                                0
-                    | Throw  -> raise <| System.Exception("You asked me to do it!")
-
-                return result
-            }
+        let cancelledToken =
+            let cts = new CancellationTokenSource()
+            cts.Cancel()
+            Some cts.Token
+
+        async { return () } |> expect Cancellation cancelledToken
 
-        let onSuccess x   =
-            match whatToDo with
-            | Cancel | Throw
-                -> Assert.Fail("Expected onSuccess but whatToDo was not Exit", [| whatToDo |])
-            | Exit
-                -> ()
+        async { failwith "computation failed" } |> expect Exception None
 
-        let onException x =
-            match whatToDo with
-            | Exit | Cancel
-                -> Assert.Fail("Expected onException but whatToDo was not Throw", [| whatToDo |])
-            | Throw  -> ()
+        async { return () } |> expect Success None
 
-        let onCancel x    =
-            match whatToDo with
-            | Exit | Throw
-                -> Assert.Fail("Expected onCancel but whatToDo was not Cancel", [| whatToDo |])
-            | Cancel -> ()
 
-        // Run it once.
-        whatToDo <- Exit
-        Async.StartWithContinuations(asyncWorkflow(), onSuccess, onException, onCancel)
 
-        whatToDo <- Cancel
-        Async.StartWithContinuations(asyncWorkflow(), onSuccess, onException, onCancel)
+type AsyncType() =
 
-        whatToDo <- Throw
-        Async.StartWithContinuations(asyncWorkflow(), onSuccess, onException, onCancel)
+    let ignoreSynchCtx f =
+        f ()
 
-        ()
+    [<VolatileField>]
+    let mutable spinloop = true
+
+    let waitASec (t:Task) =
+        let result = t.Wait(TimeSpan(hours=0,minutes=0,seconds=1))
+        Assert.True(result, "Task did not finish after waiting for a second.")
 
     [<Fact>]
     member _.AsyncRunSynchronouslyReusesThreadPoolThread() =