Pull klog bridge out to log util

argoproj · Apr 10, 2023 · 960ea46 · 960ea46
1 parent 61fd3c1
commit 960ea46
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 7 deletions.
diff --git a/cmd/kubectl-argo-rollouts/main.go b/cmd/kubectl-argo-rollouts/main.go
@@ -3,7 +3,8 @@ package main
 import (
 	"os"
 
-	"github.com/bombsimon/logrusr/v2"
+	logutil "github.com/argoproj/argo-rollouts/utils/log"
+
 	"github.com/sirupsen/logrus"
 
 	"k8s.io/cli-runtime/pkg/genericclioptions"
@@ -18,9 +19,7 @@ import (
 
 func main() {
 	klog.InitFlags(nil)
-	logrusLog := logrus.New()
-	logger := logrusr.New(logrusLog)
-	klog.SetLogger(logger)
+	logutil.SetKLogLogger(logrus.New())
 	streams := genericclioptions.IOStreams{In: os.Stdin, Out: os.Stdout, ErrOut: os.Stderr}
 	o := options.NewArgoRolloutsOptions(streams)
 	root := cmd.NewCmdArgoRollouts(o)

diff --git a/cmd/rollouts-controller/main.go b/cmd/rollouts-controller/main.go
@@ -6,6 +6,8 @@ import (
 	"strings"
 	"time"
 
+	"github.com/sirupsen/logrus"
+
 	"github.com/argoproj/pkg/kubeclientmetrics"
 	smiclientset "github.com/servicemeshinterface/smi-sdk-go/pkg/gen/client/split/clientset/versioned"
 	log "github.com/sirupsen/logrus"
@@ -83,6 +85,7 @@ func newCommand() *cobra.Command {
 			if logFormat != "" {
 				log.SetFormatter(createFormatter(logFormat))
 			}
+			logutil.SetKLogLogger(logrus.New())
 			logutil.SetKLogLevel(klogLevel)
 			log.WithField("version", version.GetVersion()).Info("Argo Rollouts starting")
 

diff --git a/utils/log/log.go b/utils/log/log.go
@@ -31,12 +31,14 @@ const (
 	NamespaceKey = "namespace"
 )
 
+// SetKLogLogger set the klog logger for the k8s go-client
+func SetKLogLogger(logger *logrus.Logger) {
+	klog.SetLogger(logrusr.New(logger))
+}
+
 // SetKLogLevel set the klog level for the k8s go-client
 func SetKLogLevel(klogLevel int) {
 	klog.InitFlags(nil)
-	logrusLog := logrus.New()
-	logger := logrusr.New(logrusLog)
-	klog.SetLogger(logger)
 	_ = flag.Set("logtostderr", "true")
 	_ = flag.Set("v", strconv.Itoa(klogLevel))
 }

diff --git a/utils/log/log_test.go b/utils/log/log_test.go
@@ -6,6 +6,8 @@ import (
 	"strings"
 	"testing"
 
+	"k8s.io/klog/v2"
+
 	log "github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/assert"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
@@ -274,3 +276,13 @@ func TestWithVersionFields(t *testing.T) {
 	assert.True(t, strings.Contains(logMessage, "generation=2"))
 	assert.True(t, strings.Contains(logMessage, "resourceVersion=123"))
 }
+
+func TestKLogLogger(t *testing.T) {
+	buf := bytes.NewBufferString("")
+	logger := log.New()
+	logger.SetOutput(buf)
+	SetKLogLogger(logger)
+	klog.Info("Logging from klog")
+	logMessage := buf.String()
+	assert.Contains(t, logMessage, "Logging from klog")
+}