From 58ce7ea568e5c0828cef796b3851201785cfc4f5 Mon Sep 17 00:00:00 2001
From: ThinkChaos <ThinkChaos@users.noreply.github.com>
Date: Thu, 4 Apr 2024 18:30:23 -0400
Subject: [PATCH] refactor(querylog): make `LoggerWriter` logger configurable

For use in tests.
---
 querylog/logger_writer.go               |  7 ++-----
 resolver/query_logging_resolver.go      |  6 ++++--
 resolver/query_logging_resolver_test.go | 12 +++++++++++-
 3 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/querylog/logger_writer.go b/querylog/logger_writer.go
index ed10874e8..fbb3cf2ea 100644
--- a/querylog/logger_writer.go
+++ b/querylog/logger_writer.go
@@ -3,19 +3,16 @@ package querylog
 import (
 	"strings"
 
-	"github.com/0xERR0R/blocky/log"
 	"github.com/0xERR0R/blocky/util"
 	"github.com/sirupsen/logrus"
 )
 
-const loggerPrefixLoggerWriter = "queryLog"
-
 type LoggerWriter struct {
 	logger *logrus.Entry
 }
 
-func NewLoggerWriter() *LoggerWriter {
-	return &LoggerWriter{logger: log.PrefixedLog(loggerPrefixLoggerWriter)}
+func NewLoggerWriter(logger *logrus.Entry) *LoggerWriter {
+	return &LoggerWriter{logger}
 }
 
 func (d *LoggerWriter) Write(entry *LogEntry) {
diff --git a/resolver/query_logging_resolver.go b/resolver/query_logging_resolver.go
index ea5493279..5de958af0 100644
--- a/resolver/query_logging_resolver.go
+++ b/resolver/query_logging_resolver.go
@@ -17,6 +17,8 @@ const (
 	cleanUpRunPeriod         = 12 * time.Hour
 	queryLoggingResolverType = "query_logging"
 	logChanCap               = 1000
+
+	loggerPrefixLoggerWriter = "queryLog"
 )
 
 // QueryLoggingResolver writes query information (question, answer, duration, ...)
@@ -50,7 +52,7 @@ func NewQueryLoggingResolver(ctx context.Context, cfg config.QueryLog) *QueryLog
 				writer, err = querylog.NewDatabaseWriter(ctx, "postgresql", cfg.Target, cfg.LogRetentionDays,
 					cfg.FlushInterval.ToDuration())
 			case config.QueryLogTypeConsole:
-				writer = querylog.NewLoggerWriter()
+				writer = querylog.NewLoggerWriter(log.PrefixedLog(loggerPrefixLoggerWriter))
 			case config.QueryLogTypeNone:
 				writer = querylog.NewNoneWriter()
 			}
@@ -68,7 +70,7 @@ func NewQueryLoggingResolver(ctx context.Context, cfg config.QueryLog) *QueryLog
 	if err != nil {
 		logger.Error("can't create query log writer, using console as fallback: ", err)
 
-		writer = querylog.NewLoggerWriter()
+		writer = querylog.NewLoggerWriter(log.PrefixedLog(loggerPrefixLoggerWriter))
 		cfg.Type = config.QueryLogTypeConsole
 	}
 
diff --git a/resolver/query_logging_resolver_test.go b/resolver/query_logging_resolver_test.go
index 4e7968a38..e46b3389f 100644
--- a/resolver/query_logging_resolver_test.go
+++ b/resolver/query_logging_resolver_test.go
@@ -13,6 +13,8 @@ import (
 	. "github.com/0xERR0R/blocky/helpertest"
 	"github.com/0xERR0R/blocky/log"
 	"github.com/0xERR0R/blocky/querylog"
+	"github.com/sirupsen/logrus"
+	"github.com/stretchr/testify/mock"
 
 	"github.com/0xERR0R/blocky/config"
 	. "github.com/0xERR0R/blocky/model"
@@ -21,7 +23,6 @@ import (
 	"github.com/miekg/dns"
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
-	"github.com/stretchr/testify/mock"
 )
 
 type SlowMockWriter struct {
@@ -44,6 +45,7 @@ var _ = Describe("QueryLoggingResolver", func() {
 		m          *mockResolver
 		tmpDir     *TmpFolder
 		mockAnswer *dns.Msg
+		writer     *log.MockLoggerHook
 
 		ctx      context.Context
 		cancelFn context.CancelFunc
@@ -69,6 +71,14 @@ var _ = Describe("QueryLoggingResolver", func() {
 		}
 
 		sut = NewQueryLoggingResolver(ctx, sutConfig)
+
+		if sutConfig.Type == config.QueryLogTypeConsole {
+			var logger *logrus.Entry
+
+			logger, writer = log.NewMockEntry()
+			sut.writer = querylog.NewLoggerWriter(logger)
+		}
+
 		m = &mockResolver{}
 		m.On("Resolve", mock.Anything).Return(&Response{Res: mockAnswer, Reason: "reason"}, nil)
 		sut.Next(m)