From 696fe0ebebb6b72ad25b598354ae5473d1126775 Mon Sep 17 00:00:00 2001
From: Vaibhav Malik <34866732+VaibhavMalik4187@users.noreply.github.com>
Date: Tue, 13 Feb 2024 15:32:49 +0530
Subject: [PATCH] Added missing tests for the go/streamlog package (#15064)

Signed-off-by: VaibhavMalik4187 <vaibhavmalik2018@gmail.com>
---
 go/streamlog/streamlog_test.go | 124 +++++++++++++++++++++++++++++++++
 1 file changed, 124 insertions(+)

diff --git a/go/streamlog/streamlog_test.go b/go/streamlog/streamlog_test.go
index 9c0b0366a1d..e653de74ea8 100644
--- a/go/streamlog/streamlog_test.go
+++ b/go/streamlog/streamlog_test.go
@@ -18,6 +18,7 @@ package streamlog
 
 import (
 	"bufio"
+	"bytes"
 	"fmt"
 	"io"
 	"net"
@@ -29,6 +30,8 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/require"
+
 	"vitess.io/vitess/go/vt/servenv"
 )
 
@@ -260,3 +263,124 @@ func TestFile(t *testing.T) {
 		t.Errorf("streamlog file: want %q got %q", want, got)
 	}
 }
+
+func TestShouldEmitLog(t *testing.T) {
+	origQueryLogFilterTag := queryLogFilterTag
+	origQueryLogRowThreshold := queryLogRowThreshold
+	defer func() {
+		SetQueryLogFilterTag(origQueryLogFilterTag)
+		SetQueryLogRowThreshold(origQueryLogRowThreshold)
+	}()
+
+	tests := []struct {
+		sql              string
+		qLogFilterTag    string
+		qLogRowThreshold uint64
+		rowsAffected     uint64
+		rowsReturned     uint64
+		ok               bool
+	}{
+		{
+			sql:              "queryLogThreshold smaller than affected and returned",
+			qLogFilterTag:    "",
+			qLogRowThreshold: 2,
+			rowsAffected:     7,
+			rowsReturned:     7,
+			ok:               true,
+		},
+		{
+			sql:              "queryLogThreshold greater than affected and returned",
+			qLogFilterTag:    "",
+			qLogRowThreshold: 27,
+			rowsAffected:     7,
+			rowsReturned:     17,
+			ok:               false,
+		},
+		{
+			sql:              "this doesn't contains queryFilterTag: TAG",
+			qLogFilterTag:    "special tag",
+			qLogRowThreshold: 10,
+			rowsAffected:     7,
+			rowsReturned:     17,
+			ok:               false,
+		},
+		{
+			sql:              "this contains queryFilterTag: TAG",
+			qLogFilterTag:    "TAG",
+			qLogRowThreshold: 0,
+			rowsAffected:     7,
+			rowsReturned:     17,
+			ok:               true,
+		},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.sql, func(t *testing.T) {
+			SetQueryLogFilterTag(tt.qLogFilterTag)
+			SetQueryLogRowThreshold(tt.qLogRowThreshold)
+
+			require.Equal(t, tt.ok, ShouldEmitLog(tt.sql, tt.rowsAffected, tt.rowsReturned))
+		})
+	}
+}
+
+func TestGetFormatter(t *testing.T) {
+	tests := []struct {
+		name           string
+		logger         *StreamLogger[string]
+		params         url.Values
+		val            any
+		expectedErr    string
+		expectedOutput string
+	}{
+		{
+			name: "unexpected value error",
+			logger: &StreamLogger[string]{
+				name: "test-logger",
+			},
+			params: url.Values{
+				"keys": []string{"key1", "key2"},
+			},
+			val:            "temp val",
+			expectedOutput: "Error: unexpected value of type string in test-logger!",
+			expectedErr:    "",
+		},
+		{
+			name: "mock formatter",
+			logger: &StreamLogger[string]{
+				name: "test-logger",
+			},
+			params: url.Values{
+				"keys": []string{"key1", "key2"},
+			},
+			val:         &mockFormatter{err: fmt.Errorf("formatter error")},
+			expectedErr: "formatter error",
+		},
+	}
+
+	for _, tt := range tests {
+		tt := tt
+		t.Run(tt.name, func(t *testing.T) {
+			var buffer bytes.Buffer
+			logFormatterFunc := GetFormatter[string](tt.logger)
+			err := logFormatterFunc(&buffer, tt.params, tt.val)
+			if tt.expectedErr == "" {
+				require.NoError(t, err)
+				require.Equal(t, tt.expectedOutput, buffer.String())
+			} else {
+				require.ErrorContains(t, err, tt.expectedErr)
+			}
+		})
+	}
+}
+
+type mockFormatter struct {
+	called bool
+	err    error
+}
+
+func (mf *mockFormatter) Logf(w io.Writer, params url.Values) error {
+	mf.called = true
+	return mf.err
+}