update(Context): add Wait and WithNewWaitGroup methods

context.go

@@ -46,6 +46,12 @@ func (c Context) Err() error {
 	return c.Context.Err()
 }
 
+// Wait runs c.WaitGroup.Wait().
+// If c.WaitGroup is not set, Wait panics.
+func (c Context) Wait() {
+	c.WaitGroup.Wait()
+}
+
 // AfterFunc arranges to call f in its own goroutine after c is done
 // (cancelled or timed out).
 // If c is already done, AfterFunc calls f immediately in its own goroutine.
@@ -139,6 +145,13 @@ func (c Context) WithTimeoutCause(timeout time.Duration, cause error) (Context,
 	return c.WithDeadlineCause(time.Now().Add(timeout), cause)
 }
 
+// WithNewWaitGroup returns a copy of c with WaitGroup field set to
+// a new [sync.WaitGroup].
+func (c Context) WithNewWaitGroup() Context {
+	c.WaitGroup = new(sync.WaitGroup)
+	return c
+}
+
 // WithWaitGroup returns a copy of c with WaitGroup field set to wg.
 func (c Context) WithWaitGroup(wg *sync.WaitGroup) Context {
 	c.WaitGroup = wg

example_test.go

@@ -3,7 +3,6 @@ package rx_test
 import (
 	"context"
 	"fmt"
-	"sync"
 	"time"
 
 	"github.com/b97tsk/rx"
@@ -84,9 +83,7 @@ func Example_blocking() {
 }
 
 func Example_waitGroup() {
-	var wg sync.WaitGroup
-
-	ctx := rx.NewContext(context.TODO()).WithWaitGroup(&wg)
+	ctx := rx.NewContext(context.TODO()).WithNewWaitGroup()
 
 	ctx.Go(func() {
 		for n := 1; n < 4; n++ {
@@ -98,7 +95,7 @@ func Example_waitGroup() {
 		}
 	})
 
-	wg.Wait()
+	ctx.Wait()
 
 	// Output:
 	// 1