Improves wording, test and conditional check

Author: jcmanzo

pipe/pipeline.go

@@ -181,12 +181,13 @@ func WithEventHandler(handler func(e *Event)) Option {
 }
 
 // WithStagePanicHandler sets a panic handler for the stages within a pipeline.
-// When a stage within the pipeline panics, the provided handler will be invoked, enabling
-// clients to capture the panic, such as for observability purposes.
+// When a pipeline stage panics, the provided handler will be invoked, allowing
+// the client to handle the panic in whatever way they see fit.
 //
 // Note:
-//   - The client is responsible for deciding whether to recover from the panic or propagate it further.
-//   - Currently, only the Function stage supports this functionality.
+//   - Only the Function stage supports this functionality.
+//   - The client is responsible for deciding whether to recover from the panic or panicking again.
+//   - If a panic handler is not set, the panic will be propagated normally.
 func WithStagePanicHandler(ph StagePanicHandler) Option {
 	return func(p *Pipeline) {
 		p.panicHandler = ph
@@ -279,10 +280,8 @@ func (p *Pipeline) Start(ctx context.Context) error {
 	}
 
 	for i, s := range p.stages {
-		if p.panicHandler != nil {
-			if phs, ok := s.(StagePanicHandlerAware); ok {
-				phs.SetPanicHandler(p.panicHandler)
-			}
+		if phs, ok := s.(StagePanicHandlerAware); ok && p.panicHandler != nil {
+			phs.SetPanicHandler(p.panicHandler)
 		}
 
 		var err error

pipe/pipeline_test.go

@@ -462,12 +462,11 @@ func TestFunction(t *testing.T) {
 	})
 
 	t.Run("panic with handler", func(t *testing.T) {
-		panicked := make(chan bool, 1)
+		panickedMessage := make(chan error, 1)
 		p := pipe.New(
 			pipe.WithDir(dir),
 			pipe.WithStagePanicHandler(func(err error) {
-				panicked <- true
-				assert.Equal(t, "this is a panic", err.Error())
+				panickedMessage <- err
 			}),
 		)
 		p.Add(
@@ -481,7 +480,7 @@ func TestFunction(t *testing.T) {
 		)
 
 		out, err := p.Output(ctx)
-		assert.True(t, <-panicked)
+		assert.Error(t, <-panickedMessage)
 		assert.Error(t, err)
 		assert.Empty(t, out)
 	})