Keep ReplayHandler when using With or WithGroup

Author: cschleiden

internal/workflowstate/replaylogger.go

@@ -24,14 +24,19 @@ func (rh *replayHandler) Handle(ctx context.Context, r slog.Record) error {
 	return rh.handler.Handle(ctx, r)
 }
 
-// WithAttrs implements slog.Handler.
 func (rh *replayHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
-	return rh.handler.WithAttrs(attrs)
+	return &replayHandler{
+		state:   rh.state,
+		handler: rh.handler.WithAttrs(attrs),
+	}
 }
 
 // WithGroup implements slog.Handler.
 func (rh *replayHandler) WithGroup(name string) slog.Handler {
-	return rh.handler.WithGroup(name)
+	return &replayHandler{
+		state:   rh.state,
+		handler: rh.handler.WithGroup(name),
+	}
 }
 
 var _ slog.Handler = (*replayHandler)(nil)

internal/workflowstate/replaylogger_test.go

@@ -0,0 +1,27 @@
+package workflowstate
+
+import (
+	"log/slog"
+	"testing"
+
+	"github.com/benbjohnson/clock"
+	"github.com/cschleiden/go-workflows/core"
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+)
+
+func Test_ReplayLogger_With(t *testing.T) {
+	i := core.NewWorkflowInstance(uuid.NewString(), "")
+	wfState := NewWorkflowState(i, slog.Default(), clock.New())
+
+	with := wfState.Logger().With(slog.String("foo", "bar"))
+	require.IsType(t, &replayHandler{}, with.Handler())
+}
+
+func Test_ReplayLogger_WithGroup(t *testing.T) {
+	i := core.NewWorkflowInstance(uuid.NewString(), "")
+	wfState := NewWorkflowState(i, slog.Default(), clock.New())
+
+	with := wfState.Logger().WithGroup("group_name")
+	require.IsType(t, &replayHandler{}, with.Handler())
+}