feat: add ability to use custom zap core

setup.go

@@ -134,13 +134,7 @@ func SetupLogging(cfg Config) {
 		newPrimaryCore = newPrimaryCore.With([]zap.Field{zap.String(k, v)})
 	}
 
-	if primaryCore != nil {
-		loggerCore.ReplaceCore(primaryCore, newPrimaryCore)
-	} else {
-		loggerCore.AddCore(newPrimaryCore)
-	}
-	primaryCore = newPrimaryCore
-
+	setPrimaryCore(newPrimaryCore)
 	setAllLoggers(defaultLevel)
 
 	for name, level := range cfg.SubsystemLevels {
@@ -152,6 +146,24 @@ func SetupLogging(cfg Config) {
 	}
 }
 
+// SetPrimaryCore changes the primary logging core. If the SetupLogging was
+// called then the previously configured core will be replaced.
+func SetPrimaryCore(core zapcore.Core) {
+	loggerMutex.Lock()
+	defer loggerMutex.Unlock()
+
+	setPrimaryCore(core)
+}
+
+func setPrimaryCore(core zapcore.Core) {
+	if primaryCore != nil {
+		loggerCore.ReplaceCore(primaryCore, core)
+	} else {
+		loggerCore.AddCore(core)
+	}
+	primaryCore = core
+}
+
 // SetDebugLogging calls SetAllLoggers with logging.DEBUG
 func SetDebugLogging() {
 	SetAllLoggers(LevelDebug)

setup_test.go

@@ -7,6 +7,8 @@ import (
 	"os"
 	"strings"
 	"testing"
+
+	"go.uber.org/zap"
 )
 
 func TestGetLoggerDefault(t *testing.T) {
@@ -36,7 +38,6 @@ func TestGetLoggerDefault(t *testing.T) {
 	if !strings.Contains(buf.String(), "scooby") {
 		t.Errorf("got %q, wanted it to contain log output", buf.String())
 	}
-
 }
 
 func TestLogToFileAndStderr(t *testing.T) {
@@ -203,3 +204,37 @@ func TestSubsystemLevels(t *testing.T) {
 		t.Errorf("got %q, wanted it to contain info2", buf.String())
 	}
 }
+
+func TestCustomCore(t *testing.T) {
+	r, w, err := os.Pipe()
+	if err != nil {
+		t.Fatalf("failed to open pipe: %v", err)
+	}
+
+	stderr := os.Stderr
+	os.Stderr = w
+	defer func() {
+		os.Stderr = stderr
+	}()
+
+	ws, _, err := zap.Open("stdout", "stderr")
+	if err != nil {
+		t.Fatalf("unable to open logging output: %v", err)
+	}
+
+	// logging should work with the custom core
+	SetPrimaryCore(newCore(PlaintextOutput, ws, LevelDebug))
+	log := getLogger("test")
+
+	log.Error("scooby")
+	w.Close()
+
+	buf := &bytes.Buffer{}
+	if _, err := io.Copy(buf, r); err != nil && err != io.ErrClosedPipe {
+		t.Fatalf("unexpected error: %v", err)
+	}
+
+	if !strings.Contains(buf.String(), "scooby") {
+		t.Errorf("got %q, wanted it to contain log output", buf.String())
+	}
+}