From f300b26592674e16f9e846b0fe348c8e6d74fd02 Mon Sep 17 00:00:00 2001
From: Nghia Tran <tcnghia@gmail.com>
Date: Tue, 18 Jun 2024 17:33:25 -0700
Subject: [PATCH] Preserve context in `With()` calls (#17)

Currently `With()` will reset context back to `context.Background()`,
causing the `wrap/wrapf` calls to lose the original context.

This affects the `clog/gcp` handler who relies on a value keyed with
`"trace"` being set on the context
([see](https://github.com/chainguard-dev/clog/blob/main/gcp/trace.go#L110-L116))
(we should use a typed value though, but that's in a different PR).

---------

Signed-off-by: Nghia Tran <tcnghia@gmail.com>
---
 logger.go      |  4 ++--
 logger_test.go | 13 +++++++++++++
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/logger.go b/logger.go
index 853a4ed..9011ffb 100644
--- a/logger.go
+++ b/logger.go
@@ -53,12 +53,12 @@ func With(args ...any) *Logger {
 
 // With calls [Logger.With] on the logger.
 func (l *Logger) With(args ...any) *Logger {
-	return NewLogger(l.Logger.With(args...))
+	return NewLoggerWithContext(l.context(), l.Logger.With(args...))
 }
 
 // WithGroup calls [Logger.WithGroup] on the default logger.
 func (l *Logger) WithGroup(name string) *Logger {
-	return NewLogger(l.Logger.WithGroup(name))
+	return NewLoggerWithContext(l.context(), l.Logger.WithGroup(name))
 }
 
 func (l *Logger) context() context.Context {
diff --git a/logger_test.go b/logger_test.go
index c48fed7..3bbc3c4 100644
--- a/logger_test.go
+++ b/logger_test.go
@@ -113,3 +113,16 @@ func TestLoggerPC(t *testing.T) {
 		t.Errorf("want %v, got %v", want, got)
 	}
 }
+
+func TestWith(t *testing.T) {
+	ctx := context.WithValue(context.Background(), "test", "test")
+	log := NewLoggerWithContext(ctx, nil)
+	withed := log.With("a", "b")
+	if want := withed.ctx; want != ctx {
+		t.Errorf("want %v, got %v", want, ctx)
+	}
+	withed = log.WithGroup("a")
+	if want := withed.ctx; want != ctx {
+		t.Errorf("want %v, got %v", want, ctx)
+	}
+}