Fix behaviour of typedCore.With

belimawr · belimawr · commit afc44c77dbb3 · 2024-04-09T19:13:47.000+02:00
diff --git a/logp/configure/logging.go b/logp/configure/logging.go
@@ -96,7 +96,7 @@ func LoggingWithTypedOutputs(beatName string, cfg, typedCfg *config.C, logKey, k
 	typedLogpConfig.Beat = beatName
 	if typedCfg != nil {
 		if err := typedCfg.Unpack(&typedLogpConfig); err != nil {
-			return fmt.Errorf("cannot unpack sensitiveCfg: %w", err)
+			return fmt.Errorf("cannot unpack typed output config: %w", err)
 		}
 	}
 
diff --git a/logp/core_test.go b/logp/core_test.go
@@ -537,24 +537,78 @@ func TestTypedLoggerCoreSync(t *testing.T) {
 }
 
 func TestTypedLoggerCoreWith(t *testing.T) {
-	defaultCoreMock := &ZapCoreMock{}
-	typedCoreMock := &ZapCoreMock{}
+	defaultWriter := writeSyncer{}
+	typedWriter := writeSyncer{}
+
+	cfg := zap.NewProductionEncoderConfig()
+	cfg.TimeKey = "" // remove the time to make the log entry consistent
+
+	defaultCore := zapcore.NewCore(
+		zapcore.NewJSONEncoder(cfg),
+		&defaultWriter,
+		zapcore.InfoLevel,
+	)
+
+	typedCore := zapcore.NewCore(
+		zapcore.NewJSONEncoder(cfg),
+		&typedWriter,
+		zapcore.InfoLevel,
+	)
 
-	defaultCoreMock.WithFunc = func(fields []zapcore.Field) zapcore.Core { return defaultCoreMock }
-	typedCoreMock.WithFunc = func(fields []zapcore.Field) zapcore.Core { return typedCoreMock }
 	core := typedLoggerCore{
-		defaultCore: defaultCoreMock,
-		typedCore:   typedCoreMock,
+		defaultCore: defaultCore,
+		typedCore:   typedCore,
+		key:         "log.type",
+		value:       "sensitive",
 	}
 
-	field := strField("foo", "bar")
-	core.With([]zapcore.Field{field})
+	expectedLines := []string{
+		// First/Default logger
+		`{"level":"info","msg":"Very first message"}`,
+
+		// Two messages after calling With
+		`{"level":"info","msg":"a message with extra fields","foo":"bar"}`,
+		`{"level":"info","msg":"another message with extra fields","foo":"bar"}`,
+
+		// A message with the default logger
+		`{"level":"info","msg":"a message without extra fields"}`,
+
+		// Two more messages with a different field
+		`{"level":"info","msg":"a message with an answer","answer":"42"}`,
+		`{"level":"info","msg":"another message with an answer","answer":"42"}`,
 
-	if core.defaultCore != defaultCoreMock {
-		t.Error("defaultCore must not change after call to With")
+		// One last message with the default logger
+		`{"level":"info","msg":"another message without any extra fields"}`,
 	}
-	if core.typedCore != typedCoreMock {
-		t.Error("typedCore must not change after call to With")
+
+	// The default logger, it should not be modified by any call to With.
+	logger := zap.New(&core)
+	logger.Info("Very first message")
+
+	// Add a field and write messages
+	loggerWithFields := logger.With(strField("foo", "bar"))
+	loggerWithFields.Info("a message with extra fields")
+	loggerWithFields.Info("another message with extra fields")
+
+	// Use the default logger again
+	logger.Info("a message without extra fields")
+
+	// New logger with other fields
+	loggerWithFields = logger.With(strField("answer", "42"))
+	loggerWithFields.Info("a message with an answer")
+	loggerWithFields.Info("another message with an answer")
+
+	// One last message with the default logger
+	logger.Info("another message without any extra fields")
+
+	scanner := bufio.NewScanner(strings.NewReader(defaultWriter.String()))
+	count := 0
+	for scanner.Scan() {
+		l := scanner.Text()
+		if l != expectedLines[count] {
+			t.Error("Expecting:\n", l, "\nGot:\n", expectedLines[count])
+		}
+		count++
 	}
 }
 
diff --git a/logp/typedloggercore.go b/logp/typedloggercore.go
@@ -43,9 +43,13 @@ func (t *typedLoggerCore) Enabled(l zapcore.Level) bool {
 }
 
 func (t *typedLoggerCore) With(fields []zapcore.Field) zapcore.Core {
-	t.defaultCore = t.defaultCore.With(fields)
-	t.typedCore = t.typedCore.With(fields)
-	return t
+	newCore := typedLoggerCore{
+		defaultCore: t.defaultCore.With(fields),
+		typedCore:   t.typedCore.With(fields),
+		key:         t.key,
+		value:       t.value,
+	}
+	return &newCore
 }
 
 func (t *typedLoggerCore) Check(e zapcore.Entry, ce *zapcore.CheckedEntry) *zapcore.CheckedEntry {

Original file line number	Diff line number	Diff line change
`@@ -96,7 +96,7 @@ func LoggingWithTypedOutputs(beatName string, cfg, typedCfg *config.C, logKey, k`
`96`	`96`	`typedLogpConfig.Beat = beatName`
`97`	`97`	`if typedCfg != nil {`
`98`	`98`	`if err := typedCfg.Unpack(&typedLogpConfig); err != nil {`
`99`		`- return fmt.Errorf("cannot unpack sensitiveCfg: %w", err)`
	`99`	`+ return fmt.Errorf("cannot unpack typed output config: %w", err)`
`100`	`100`	`}`
`101`	`101`	`}`
`102`	`102`