From 925cf8d3d1e76470ae84f8af2330c391d7202952 Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Fri, 19 May 2023 12:15:43 +0800
Subject: [PATCH] chore: add more tests (#3261)

---
 core/prof/runtime.go      |  8 +++++++-
 core/prof/runtime_test.go | 36 ++++++++++++++++++++++++++++++++++++
 2 files changed, 43 insertions(+), 1 deletion(-)
 create mode 100644 core/prof/runtime_test.go

diff --git a/core/prof/runtime.go b/core/prof/runtime.go
index 17ede442..af41d051 100644
--- a/core/prof/runtime.go
+++ b/core/prof/runtime.go
@@ -2,6 +2,8 @@ package prof
 
 import (
 	"fmt"
+	"io"
+	"os"
 	"runtime"
 	"time"
 )
@@ -13,6 +15,10 @@ const (
 
 // DisplayStats prints the goroutine, memory, GC stats with given interval, default to 5 seconds.
 func DisplayStats(interval ...time.Duration) {
+	displayStatsWithWriter(os.Stdout, interval...)
+}
+
+func displayStatsWithWriter(writer io.Writer, interval ...time.Duration) {
 	duration := defaultInterval
 	for _, val := range interval {
 		duration = val
@@ -24,7 +30,7 @@ func DisplayStats(interval ...time.Duration) {
 		for range ticker.C {
 			var m runtime.MemStats
 			runtime.ReadMemStats(&m)
-			fmt.Printf("Goroutines: %d, Alloc: %vm, TotalAlloc: %vm, Sys: %vm, NumGC: %v\n",
+			fmt.Fprintf(writer, "Goroutines: %d, Alloc: %vm, TotalAlloc: %vm, Sys: %vm, NumGC: %v\n",
 				runtime.NumGoroutine(), m.Alloc/mega, m.TotalAlloc/mega, m.Sys/mega, m.NumGC)
 		}
 	}()
diff --git a/core/prof/runtime_test.go b/core/prof/runtime_test.go
new file mode 100644
index 00000000..208e75f5
--- /dev/null
+++ b/core/prof/runtime_test.go
@@ -0,0 +1,36 @@
+package prof
+
+import (
+	"strings"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestDisplayStats(t *testing.T) {
+	writer := &threadSafeBuffer{
+		buf: strings.Builder{},
+	}
+	displayStatsWithWriter(writer, time.Millisecond*10)
+	time.Sleep(time.Millisecond * 50)
+	assert.Contains(t, writer.String(), "Goroutines: ")
+}
+
+type threadSafeBuffer struct {
+	buf  strings.Builder
+	lock sync.Mutex
+}
+
+func (b *threadSafeBuffer) String() string {
+	b.lock.Lock()
+	defer b.lock.Unlock()
+	return b.buf.String()
+}
+
+func (b *threadSafeBuffer) Write(p []byte) (n int, err error) {
+	b.lock.Lock()
+	defer b.lock.Unlock()
+	return b.buf.Write(p)
+}