From 80a0096ee2bdc8f26a95489333943ae04b5a6244 Mon Sep 17 00:00:00 2001 From: zhifei92 Date: Wed, 23 Oct 2024 22:30:32 +0800 Subject: [PATCH] fix: Add a thread-safe logBuffer for the watchdog unit test. --- pkg/kubelet/watchdog/watchdog_linux_test.go | 45 +++++++++++++++++---- 1 file changed, 37 insertions(+), 8 deletions(-) diff --git a/pkg/kubelet/watchdog/watchdog_linux_test.go b/pkg/kubelet/watchdog/watchdog_linux_test.go index 86cdd744cd4..ed167742523 100644 --- a/pkg/kubelet/watchdog/watchdog_linux_test.go +++ b/pkg/kubelet/watchdog/watchdog_linux_test.go @@ -25,6 +25,7 @@ import ( "flag" "net/http" "strings" + "sync" "testing" "time" @@ -139,14 +140,7 @@ func TestHealthCheckerStart(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { // Capture logs - var logBuffer bytes.Buffer - flags := &flag.FlagSet{} - klog.InitFlags(flags) - if err := flags.Set("v", "5"); err != nil { - t.Fatal(err) - } - klog.LogToStderr(false) - klog.SetOutput(&logBuffer) + logBuffer := setupLogging(t) // Mock SdWatchdogEnabled to return a valid value mockClient := &mockWatchdogClient{ @@ -178,3 +172,38 @@ func TestHealthCheckerStart(t *testing.T) { }) } } + +// threadSafeBuffer is a thread-safe wrapper around bytes.Buffer. +type threadSafeBuffer struct { + buffer bytes.Buffer + mu sync.Mutex +} + +func (b *threadSafeBuffer) Write(p []byte) (n int, err error) { + b.mu.Lock() + defer b.mu.Unlock() + return b.buffer.Write(p) +} + +func (b *threadSafeBuffer) String() string { + b.mu.Lock() + defer b.mu.Unlock() + return b.buffer.String() +} + +// setupLogging sets up logging to capture output using a thread-safe buffer. +func setupLogging(t *testing.T) *threadSafeBuffer { + flags := &flag.FlagSet{} + klog.InitFlags(flags) + if err := flags.Set("v", "5"); err != nil { + t.Fatal(err) + } + klog.LogToStderr(false) + + logBuffer := &threadSafeBuffer{} + + // Set the output to the thread-safe buffer + klog.SetOutput(logBuffer) + + return logBuffer +}