fix(log): fix caller skip depth

reearth · Jan 28, 2025 · c3eb734 · c3eb734
1 parent 83ed95f
commit c3eb734
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 10 deletions.
diff --git a/log/context.go b/log/context.go
@@ -20,9 +20,13 @@ func GetLoggerFromContext(ctx context.Context) *Logger {
 
 func GetLoggerFromContextOrDefault(ctx context.Context) *Logger {
 	if logger := GetLoggerFromContext(ctx); logger != nil {
-		return logger
+		return logger.AddCallerSkip(1)
 	}
-	return globalLogger
+	return globalLogger.AddCallerSkip(1)
+}
+
+func UpdateContext(ctx context.Context, f func(logger *Logger) *Logger) context.Context {
+	return AttachLoggerToContext(ctx, f(GetLoggerFromContextOrDefault(ctx)))
 }
 
 func Tracefc(ctx context.Context, format string, args ...any) {

diff --git a/log/context_test.go b/log/context_test.go
@@ -16,21 +16,22 @@ func TestContextLogger(t *testing.T) {
 		SetOutput(DefaultOutput)
 	})
 
-	l := NewWithOutput(w).SetPrefix("test")
+	l := NewWithOutput(w).SetPrefix("prefix")
 	ctx := AttachLoggerToContext(context.Background(), l)
+
 	Infofc(ctx, "hoge %s", "fuga")
 	Infofc(context.Background(), "hoge %s", "fuga2")
 	//nolint:staticcheck // test nil context
 	Infofc(nil, "hoge %s", "fuga3")
 
 	scanner := bufio.NewScanner(w)
 	assert.True(t, scanner.Scan())
-	assert.Contains(t, scanner.Text(), "test\thoge fuga")
+	assert.Contains(t, scanner.Text(), "\tprefix\t")
 	assert.True(t, scanner.Scan())
 	assert.Contains(t, scanner.Text(), "hoge fuga2")
-	assert.NotContains(t, scanner.Text(), "test")
+	assert.NotContains(t, scanner.Text(), "\tprefix\t")
 	assert.True(t, scanner.Scan())
 	assert.Contains(t, scanner.Text(), "hoge fuga3")
-	assert.NotContains(t, scanner.Text(), "test")
+	assert.NotContains(t, scanner.Text(), "\tprefix\t")
 	assert.False(t, scanner.Scan())
 }
diff --git a/log/echo_test.go b/log/echo_test.go
@@ -12,11 +12,11 @@ func TestEcho(t *testing.T) {
 	w := &bytes.Buffer{}
 	l := NewEcho()
 	l.SetOutput(w)
-	l.SetPrefix("test")
+	l.SetPrefix("prefix")
 	l.Infof("hoge %s", "fuga")
 
 	scanner := bufio.NewScanner(w)
 	assert.True(t, scanner.Scan())
-	assert.Contains(t, scanner.Text(), "test\thoge fuga")
+	assert.Contains(t, scanner.Text(), "\tprefix\t")
 	assert.False(t, scanner.Scan())
 }
diff --git a/log/logger.go b/log/logger.go
@@ -177,6 +177,26 @@ func (l *Logger) ClearPrefix() *Logger {
 	}
 }
 
+func (l *Logger) WithCaller(enabled bool) *Logger {
+	return &Logger{
+		logger:    l.logger.WithOptions(zap.WithCaller(enabled)),
+		atom:      l.atom,
+		prefix:    l.prefix,
+		dynPrefix: l.dynPrefix,
+		dynSuffix: l.dynSuffix,
+	}
+}
+
+func (l *Logger) AddCallerSkip(skip int) *Logger {
+	return &Logger{
+		logger:    l.logger.WithOptions(zap.AddCallerSkip(skip)),
+		atom:      l.atom,
+		prefix:    l.prefix,
+		dynPrefix: l.dynPrefix,
+		dynSuffix: l.dynSuffix,
+	}
+}
+
 func (l *Logger) Debugf(format string, args ...any) {
 	f := l.format(format, args...)
 	l.logger.Debugf(f.Format, f.Args...)

diff --git a/log/logger_test.go b/log/logger_test.go
@@ -37,9 +37,9 @@ func TestLogger_SetPrefix(t *testing.T) {
 
 	scanner := bufio.NewScanner(w)
 	assert.True(t, scanner.Scan())
-	assert.Contains(t, scanner.Text(), "test\thoge fuga")
+	assert.Regexp(t, `\ttest\t.+?\thoge fuga$`, scanner.Text()) // .+? is a caller
 	assert.True(t, scanner.Scan())
-	assert.Contains(t, scanner.Text(), "test\t[fuga 1]")
+	assert.Regexp(t, `\ttest\t.+?\t\[fuga 1\]$`, scanner.Text())
 	assert.True(t, scanner.Scan())
 	assert.Contains(t, scanner.Text(), "test")
 	assert.Contains(t, scanner.Text(), "abcd")