From 264fcab7c6dfee0bbcee45a869f36a460cdafa95 Mon Sep 17 00:00:00 2001
From: Prashant Varanasi <prashant@uber.com>
Date: Sat, 11 Apr 2020 15:51:18 -0700
Subject: [PATCH] Fix IncreaseLevel being reset after With (#812)

Fixes #810

We missed implementing the increase-level logic in With
and since Core was embedded, we wended up using With from
the original core.

To avoid this sort of issue, avoid embedding and implement all
Core methods explicitly. This lets us consider the behaviour
of each method explicitly.
---
 increase_level_test.go         | 13 +++++++++++--
 zapcore/increase_level.go      | 17 ++++++++++++++---
 zapcore/increase_level_test.go | 25 +++++++++++++++++++++++--
 3 files changed, 48 insertions(+), 7 deletions(-)

diff --git a/increase_level_test.go b/increase_level_test.go
index d3e111464..56a23873c 100644
--- a/increase_level_test.go
+++ b/increase_level_test.go
@@ -29,10 +29,13 @@ import (
 	"go.uber.org/zap/zaptest/observer"
 )
 
-func newLoggedEntry(level zapcore.Level, msg string) observer.LoggedEntry {
+func newLoggedEntry(level zapcore.Level, msg string, fields ...zapcore.Field) observer.LoggedEntry {
+	if len(fields) == 0 {
+		fields = []zapcore.Field{}
+	}
 	return observer.LoggedEntry{
 		Entry:   zapcore.Entry{Level: level, Message: msg},
-		Context: []zapcore.Field{},
+		Context: fields,
 	}
 }
 
@@ -75,9 +78,15 @@ func TestIncreaseLevel(t *testing.T) {
 		errorLogger.Warn("ignored warn log")
 		errorLogger.Error("increase level error log")
 
+		withFields := errorLogger.With(String("k", "v"))
+		withFields.Debug("ignored debug log with fields")
+		withFields.Warn("ignored warn log with fields")
+		withFields.Error("increase level error log with fields")
+
 		assert.Equal(t, []observer.LoggedEntry{
 			newLoggedEntry(WarnLevel, "original warn log"),
 			newLoggedEntry(ErrorLevel, "increase level error log"),
+			newLoggedEntry(ErrorLevel, "increase level error log with fields", String("k", "v")),
 		}, logs.AllUntimed(), "unexpected logs")
 
 		assert.Empty(t, errorOut.String(), "expect no error output")
diff --git a/zapcore/increase_level.go b/zapcore/increase_level.go
index a42135c15..5a1749261 100644
--- a/zapcore/increase_level.go
+++ b/zapcore/increase_level.go
@@ -23,8 +23,7 @@ package zapcore
 import "fmt"
 
 type levelFilterCore struct {
-	Core
-
+	core  Core
 	level LevelEnabler
 }
 
@@ -46,10 +45,22 @@ func (c *levelFilterCore) Enabled(lvl Level) bool {
 	return c.level.Enabled(lvl)
 }
 
+func (c *levelFilterCore) With(fields []Field) Core {
+	return &levelFilterCore{c.core.With(fields), c.level}
+}
+
 func (c *levelFilterCore) Check(ent Entry, ce *CheckedEntry) *CheckedEntry {
 	if !c.Enabled(ent.Level) {
 		return ce
 	}
 
-	return c.Core.Check(ent, ce)
+	return c.core.Check(ent, ce)
+}
+
+func (c *levelFilterCore) Write(ent Entry, fields []Field) error {
+	return c.core.Write(ent, fields)
+}
+
+func (c *levelFilterCore) Sync() error {
+	return c.core.Sync()
 }
diff --git a/zapcore/increase_level_test.go b/zapcore/increase_level_test.go
index 8b8613375..acb8700f7 100644
--- a/zapcore/increase_level_test.go
+++ b/zapcore/increase_level_test.go
@@ -26,6 +26,7 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/zap"
 	. "go.uber.org/zap/zapcore"
 	"go.uber.org/zap/zaptest/observer"
 )
@@ -35,6 +36,7 @@ func TestIncreaseLevel(t *testing.T) {
 		coreLevel     Level
 		increaseLevel Level
 		wantErr       bool
+		with          []Field
 	}{
 		{
 			coreLevel:     InfoLevel,
@@ -49,6 +51,11 @@ func TestIncreaseLevel(t *testing.T) {
 			coreLevel:     InfoLevel,
 			increaseLevel: ErrorLevel,
 		},
+		{
+			coreLevel:     InfoLevel,
+			increaseLevel: ErrorLevel,
+			with:          []Field{zap.String("k", "v")},
+		},
 		{
 			coreLevel:     ErrorLevel,
 			increaseLevel: DebugLevel,
@@ -73,7 +80,7 @@ func TestIncreaseLevel(t *testing.T) {
 	for _, tt := range tests {
 		msg := fmt.Sprintf("increase %v to %v", tt.coreLevel, tt.increaseLevel)
 		t.Run(msg, func(t *testing.T) {
-			logger, _ := observer.New(tt.coreLevel)
+			logger, logs := observer.New(tt.coreLevel)
 
 			filteredLogger, err := NewIncreaseLevelCore(logger, tt.increaseLevel)
 			if tt.wantErr {
@@ -82,19 +89,33 @@ func TestIncreaseLevel(t *testing.T) {
 				return
 			}
 
+			if len(tt.with) > 0 {
+				filteredLogger = filteredLogger.With(tt.with)
+			}
+
 			require.NoError(t, err)
 
 			for l := DebugLevel; l <= FatalLevel; l++ {
 				enabled := filteredLogger.Enabled(l)
-				ce := filteredLogger.Check(Entry{Level: l}, nil)
+				entry := Entry{Level: l}
+				ce := filteredLogger.Check(entry, nil)
+				ce.Write()
+				entries := logs.TakeAll()
 
 				if l >= tt.increaseLevel {
 					assert.True(t, enabled, "expect %v to be enabled", l)
 					assert.NotNil(t, ce, "expect non-nil Check")
+					assert.NotEmpty(t, entries, "Expect log to be written")
 				} else {
 					assert.False(t, enabled, "expect %v to be disabled", l)
 					assert.Nil(t, ce, "expect nil Check")
+					assert.Empty(t, entries, "No logs should have been written")
 				}
+
+				// Write should always log the entry as per the Core interface
+				require.NoError(t, filteredLogger.Write(entry, nil), "Write failed")
+				require.NoError(t, filteredLogger.Sync(), "Sync failed")
+				assert.NotEmpty(t, logs.TakeAll(), "Write should always log")
 			}
 		})
 	}