tools/flow: fix test race

The test would sometimes not generate the panic as Value could race.

Signed-off-by: Marcel van Lohuizen <mpvl@golang.org>
Change-Id: I215cbfec66c21229b8a4a968e141dec1b46ba507
Reviewed-on: https://review.gerrithub.io/c/cue-lang/cue/+/535802
Unity-Result: CUEcueckoo <cueckoo@cuelang.org>
Reviewed-by: Paul Jolly <paul@myitcv.io>
TryBot-Result: CUEcueckoo <cueckoo@cuelang.org>
diff --git a/tools/flow/flow_test.go b/tools/flow/flow_test.go
index 48abad9..59e8ca2 100644
--- a/tools/flow/flow_test.go
+++ b/tools/flow/flow_test.go
@@ -101,30 +101,38 @@
             $id: "slow"
             out: string
         }
+        b: {
+            $id:    "slow"
+            $after: a
+            out:    string
+        }
     }
     `
 	ctx := cuecontext.New()
 	v := ctx.CompileString(f)
 
-	start := make(chan bool, 1)
+	ch := make(chan bool, 1)
 
-	cfg := &flow.Config{Root: cue.ParsePath("root")}
+	cfg := &flow.Config{
+		Root: cue.ParsePath("root"),
+		UpdateFunc: func(c *flow.Controller, t *flow.Task) error {
+			ch <- true
+			return nil
+		},
+	}
+
 	c := flow.New(cfg, v, taskFunc)
 
 	defer func() { recover() }()
 
-	go func() {
-		start <- true
-		c.Run(context.TODO())
-	}()
+	go c.Run(context.TODO())
 
-	<-start
-
-	// Wait for the flow to be running
-	// The task sleeps for 10 Milliseconds
-	time.Sleep(5 * time.Millisecond)
-	// Should trigger a panic as the flow is not terminated
+	// Call Value amidst two task runs. This should trigger a panic as the flow
+	// is not terminated.
+	<-ch
 	c.Value()
+	<-ch
+
 	t.Errorf("Value() did not panic")
 }