From ccbaaef1ab583752a331d3fafb44b467521b2b97 Mon Sep 17 00:00:00 2001
From: Anderson Queiroz <anderson.queiroz@elastic.co>
Date: Thu, 11 Apr 2024 17:06:01 +0200
Subject: [PATCH] logp: add an in-memory logger (#191)

---
 logp/logger.go      | 30 ++++++++++++++++++++++++++++++
 logp/logger_test.go | 29 +++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+)

diff --git a/logp/logger.go b/logp/logger.go
index 32fd8782..88ee1752 100644
--- a/logp/logger.go
+++ b/logp/logger.go
@@ -18,8 +18,10 @@
 package logp
 
 import (
+	"bytes"
 	"fmt"
 
+	"go.elastic.co/ecszap"
 	"go.uber.org/zap"
 	"go.uber.org/zap/zapcore"
 )
@@ -50,6 +52,34 @@ func NewLogger(selector string, options ...LogOption) *Logger {
 	return newLogger(loadLogger().rootLogger, selector, options...)
 }
 
+// NewInMemory returns a new in-memory logger along with the buffer to which it
+// logs. It's goroutine safe, but operating directly on the returned buffer is not.
+// This logger is primary intended for short and simple use-cases such as printing
+// the full logs only when an operation fails.
+// encCfg configures the log format, use logp.ConsoleEncoderConfig for console
+// format, logp.JSONEncoderConfig for JSON or any other valid zapcore.EncoderConfig.
+func NewInMemory(selector string, encCfg zapcore.EncoderConfig) (*Logger, *bytes.Buffer) {
+	buff := bytes.Buffer{}
+
+	encoderConfig := ecszap.ECSCompatibleEncoderConfig(encCfg)
+	encoderConfig.EncodeTime = zapcore.ISO8601TimeEncoder
+	encoder := zapcore.NewConsoleEncoder(encoderConfig)
+
+	core := zapcore.NewCore(
+		encoder,
+		zapcore.Lock(zapcore.AddSync(&buff)),
+		zap.NewAtomicLevelAt(zap.DebugLevel))
+	ecszap.ECSCompatibleEncoderConfig(ConsoleEncoderConfig())
+
+	logger := NewLogger(
+		selector,
+		zap.WrapCore(func(in zapcore.Core) zapcore.Core {
+			return core
+		}))
+
+	return logger, &buff
+}
+
 // WithOptions returns a clone of l with options applied.
 func (l *Logger) WithOptions(options ...LogOption) *Logger {
 	cloned := l.logger.WithOptions(options...)
diff --git a/logp/logger_test.go b/logp/logger_test.go
index eaf8a107..256ff92f 100644
--- a/logp/logger_test.go
+++ b/logp/logger_test.go
@@ -18,6 +18,7 @@
 package logp
 
 import (
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -50,3 +51,31 @@ func TestLoggerWithOptions(t *testing.T) {
 	require.Len(t, observedEntries2, 1)
 	assert.Equal(t, "hello logger1 and logger2", observedEntries2[0].Message)
 }
+
+func TestNewInMemory(t *testing.T) {
+	log, buff := NewInMemory("in_memory", ConsoleEncoderConfig())
+
+	log.Debugw("a debug message", "debug_key", "debug_val")
+	log.Infow("a info message", "info_key", "info_val")
+	log.Warnw("a warn message", "warn_key", "warn_val")
+	log.Errorw("an error message", "error_key", "error_val")
+
+	logs := strings.Split(strings.TrimSpace(buff.String()), "\n")
+	assert.Len(t, logs, 4, "expected 4 log entries")
+
+	assert.Contains(t, logs[0], "a debug message")
+	assert.Contains(t, logs[0], "debug_key")
+	assert.Contains(t, logs[0], "debug_val")
+
+	assert.Contains(t, logs[1], "a info message")
+	assert.Contains(t, logs[1], "info_key")
+	assert.Contains(t, logs[1], "info_val")
+
+	assert.Contains(t, logs[2], "a warn message")
+	assert.Contains(t, logs[2], "warn_key")
+	assert.Contains(t, logs[2], "warn_val")
+
+	assert.Contains(t, logs[3], "an error message")
+	assert.Contains(t, logs[3], "error_key")
+	assert.Contains(t, logs[3], "error_val")
+}