From 726ba4174af67ebe8705570f20598e24e8d58c28 Mon Sep 17 00:00:00 2001
From: aya <ayahassan2877@gmail.com>
Date: Thu, 13 Mar 2025 12:36:25 +0200
Subject: [PATCH] Modify logging

---
 waku/logging.go     | 46 ++++++++++++++++++++++++---------------------
 waku/stress_test.go | 34 +++++++++++++++++++++++----------
 2 files changed, 49 insertions(+), 31 deletions(-)

diff --git a/waku/logging.go b/waku/logging.go
index 6d610aa..8be327c 100644
--- a/waku/logging.go
+++ b/waku/logging.go
@@ -3,41 +3,45 @@ package waku
 import (
 	"sync"
 
-	"github.com/sirupsen/logrus"
+	"go.uber.org/zap"
 )
 
 var (
-	once     sync.Once
-	instance *logrus.Logger
+	once  sync.Once
+	sugar *zap.SugaredLogger
 )
 
-// _getLogger ensures we always return the same logger instance (private function)
-func _getLogger() *logrus.Logger {
+func _getLogger() *zap.SugaredLogger {
 	once.Do(func() {
-		instance = logrus.New()
-		instance.SetFormatter(&logrus.TextFormatter{
-			FullTimestamp: true,
-		})
-		instance.SetLevel(logrus.DebugLevel) // Set default log level
+
+		config := zap.NewDevelopmentConfig()
+		l, err := config.Build()
+		if err != nil {
+			panic(err)
+		}
+		sugar = l.Sugar()
 	})
-	return instance
+	return sugar
+}
+
+func SetLogger(newLogger *zap.Logger) {
+	once.Do(func() {})
+
+	sugar = newLogger.Sugar()
 }
 
-// Debug logs a debug message
 func Debug(msg string, args ...interface{}) {
-	_getLogger().WithFields(logrus.Fields{}).Debugf(msg, args...)
+	_getLogger().Debugf(msg, args...)
 }
 
-// Info logs an info message
 func Info(msg string, args ...interface{}) {
-	_getLogger().WithFields(logrus.Fields{}).Infof(msg, args...)
-}
-
-// Error logs an error message
-func Error(msg string, args ...interface{}) {
-	_getLogger().WithFields(logrus.Fields{}).Errorf(msg, args...)
+	_getLogger().Infof(msg, args...)
 }
 
 func Warn(msg string, args ...interface{}) {
-	_getLogger().WithFields(logrus.Fields{}).Warnf(msg, args...)
+	_getLogger().Warnf(msg, args...)
+}
+
+func Error(msg string, args ...interface{}) {
+	_getLogger().Errorf(msg, args...)
 }
diff --git a/waku/stress_test.go b/waku/stress_test.go
index c357ad6..5dcb053 100644
--- a/waku/stress_test.go
+++ b/waku/stress_test.go
@@ -3,15 +3,15 @@ package waku
 import (
 	"fmt"
 	"io"
-	"log"
 	"os"
 	"runtime"
 	"testing"
 	"time"
 
-	"github.com/sirupsen/logrus"
 	"github.com/stretchr/testify/require"
 	"github.com/waku-org/waku-go-bindings/waku/common"
+	"go.uber.org/zap"
+	"go.uber.org/zap/zapcore"
 	"google.golang.org/protobuf/proto"
 )
 
@@ -21,10 +21,17 @@ func TestMemoryUsageForThreeNodes(t *testing.T) {
 	defer logFile.Close()
 
 	multiWriter := io.MultiWriter(os.Stdout, logFile)
-	log.SetOutput(multiWriter)
-	logger := _getLogger()
-	logger.SetOutput(multiWriter)
-	logger.SetLevel(logrus.DebugLevel)
+	encoderCfg := zap.NewProductionEncoderConfig()
+	encoderCfg.TimeKey = "ts"
+	encoderCfg.EncodeTime = zapcore.ISO8601TimeEncoder
+
+	core := zapcore.NewCore(
+		zapcore.NewConsoleEncoder(encoderCfg),
+		zapcore.AddSync(multiWriter),
+		zap.DebugLevel,
+	)
+	customLogger := zap.New(core)
+	SetLogger(customLogger)
 
 	testName := t.Name()
 
@@ -85,9 +92,16 @@ func Test2Nodes500IterationTearDown(t *testing.T) {
 	defer logFile.Close()
 
 	multiWriter := io.MultiWriter(os.Stdout, logFile)
-	logger := _getLogger()
-	logger.SetOutput(multiWriter)
-	logger.SetLevel(logrus.DebugLevel)
+	encoderCfg := zap.NewProductionEncoderConfig()
+	encoderCfg.TimeKey = "ts"
+	encoderCfg.EncodeTime = zapcore.ISO8601TimeEncoder
+	core := zapcore.NewCore(
+		zapcore.NewConsoleEncoder(encoderCfg),
+		zapcore.AddSync(multiWriter),
+		zap.DebugLevel,
+	)
+	customLogger := zap.New(core)
+	SetLogger(customLogger)
 
 	var memStats runtime.MemStats
 	runtime.ReadMemStats(&memStats)
@@ -210,5 +224,5 @@ func TestStoreQuery5kMessagesWithPagination(t *testing.T) {
 
 	require.LessOrEqual(t, finalHeapAlloc, initialHeapAlloc*2, "Memory usage has grown too much")
 
-	Debug("Test completed successfully")
+	Debug("[%s] Test completed successfully", t.Name())
 }