From 19b7e51bc1d43e775bfed5be6c61134ed3567682 Mon Sep 17 00:00:00 2001
From: b97tsk <b97tsk@gmail.com>
Date: Sun, 10 Mar 2024 08:00:00 +0800
Subject: [PATCH] update(Context): add Wait and WithNewWaitGroup methods

---
 context.go      | 13 +++++++++++++
 example_test.go |  7 ++-----
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/context.go b/context.go
index 552d09e9..17bd4946 100644
--- a/context.go
+++ b/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
diff --git a/example_test.go b/example_test.go
index 5f1918c0..4d896452 100644
--- a/example_test.go
+++ b/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