helper/signalwrapper: more tests

helper/signalwrapper/wrapper_test.go

@@ -66,3 +66,29 @@ func TestWrapped_cancel(t *testing.T) {
 		}
 	}
 }
+
+func TestWrapped_waitAndCancel(t *testing.T) {
+	errVal := errors.New("hi")
+	readyCh := make(chan struct{})
+	f := func(ch <-chan struct{}) error {
+		<-ch
+		<-readyCh
+		return errVal
+	}
+
+	wrapped := Run(f)
+
+	// Run both cancel and wait and wait some time to hope they're
+	// scheduled. We could _ensure_ both are scheduled by using some
+	// more lines of code but this is probably just good enough.
+	go wrapped.Cancel()
+	go wrapped.Wait()
+	close(readyCh)
+	time.Sleep(10 * time.Millisecond)
+
+	// Check it by calling Cancel again
+	err := wrapped.Cancel()
+	if err != errVal {
+		t.Fatalf("bad: %#v", err)
+	}
+}