From 186c11b2cd41593d55a4610d23dca48ab0a80be8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Vin=C3=ADcius=20Garcia?= <vingarcia00@gmail.com>
Date: Tue, 11 Jul 2023 23:07:50 -0300
Subject: [PATCH] Add more tests for the logger feature

---
 ksql_test.go | 68 ++++++++++++++++++++++++++++++++++++++++++++++------
 logger.go    |  7 +++++-
 2 files changed, 67 insertions(+), 8 deletions(-)

diff --git a/ksql_test.go b/ksql_test.go
index 13105ca..cedfa28 100644
--- a/ksql_test.go
+++ b/ksql_test.go
@@ -121,10 +121,10 @@ func TestInjectLogger(t *testing.T) {
 				var row []struct {
 					Count int `ksql:"count"`
 				}
-				return db.Query(ctx, &row, `SELECT count(*) AS count FROM users WHERE type = $1 AND age < $2`, "fakeType", 42)
+				return db.Query(ctx, &row, `FROM users WHERE type = $1 AND age < $2`, "fakeType", 42)
 			},
 
-			expectLoggedQueryToContain: []string{"count(*)", "type = $1"},
+			expectLoggedQueryToContain: []string{"SELECT", "count", "type = $1"},
 			expectLoggedParams:         []interface{}{"fakeType", 42},
 		},
 		{
@@ -134,11 +134,11 @@ func TestInjectLogger(t *testing.T) {
 				var row []struct {
 					Count int `ksql:"count"`
 				}
-				return db.Query(ctx, &row, `SELECT count(*) AS count FROM users WHERE type = $1 AND age < $2`, "fakeType", 42)
+				return db.Query(ctx, &row, `FROM users WHERE type = $1 AND age < $2`, "fakeType", 42)
 			},
 			queryErr: errors.New("fakeErrMsg"),
 
-			expectLoggedQueryToContain: []string{"count(*)", "type = $1"},
+			expectLoggedQueryToContain: []string{"SELECT", "count", "type = $1"},
 			expectLoggedParams:         []interface{}{"fakeType", 42},
 			expectLoggedErrToContain:   []string{"fakeErrMsg"},
 		},
@@ -149,11 +149,54 @@ func TestInjectLogger(t *testing.T) {
 				var row []struct {
 					Count int `ksql:"count"`
 				}
-				return db.Query(ctx, &row, `SELECT count(*) AS count FROM users WHERE type = $1 AND age < $2`, "fakeType", 42)
+				return db.Query(ctx, &row, `FROM users WHERE type = $1 AND age < $2`, "fakeType", 42)
 			},
 			queryErr: errors.New("fakeErrMsg"),
 
-			expectLoggedQueryToContain: []string{"count(*)", "type = $1"},
+			expectLoggedQueryToContain: []string{"SELECT", "count", "type = $1"},
+			expectLoggedParams:         []interface{}{"fakeType", 42},
+			expectLoggedErrToContain:   []string{"fakeErrMsg"},
+		},
+		{
+			desc:     "should work for the QueryOne function",
+			logLevel: "info",
+			methodCall: func(ctx context.Context, db Provider) error {
+				var row struct {
+					Count int `ksql:"count"`
+				}
+				return db.QueryOne(ctx, &row, `FROM users WHERE type = $1 AND age < $2`, "fakeType", 42)
+			},
+
+			expectLoggedQueryToContain: []string{"SELECT", "count", "type = $1"},
+			expectLoggedParams:         []interface{}{"fakeType", 42},
+		},
+		{
+			desc:     "should work for the QueryOne function when an error is returned",
+			logLevel: "info",
+			methodCall: func(ctx context.Context, db Provider) error {
+				var row struct {
+					Count int `ksql:"count"`
+				}
+				return db.QueryOne(ctx, &row, `FROM users WHERE type = $1 AND age < $2`, "fakeType", 42)
+			},
+			queryErr: errors.New("fakeErrMsg"),
+
+			expectLoggedQueryToContain: []string{"SELECT", "count", "type = $1"},
+			expectLoggedParams:         []interface{}{"fakeType", 42},
+			expectLoggedErrToContain:   []string{"fakeErrMsg"},
+		},
+		{
+			desc:     "should work for the QueryOne function when an error is returned with error level",
+			logLevel: "error",
+			methodCall: func(ctx context.Context, db Provider) error {
+				var row struct {
+					Count int `ksql:"count"`
+				}
+				return db.QueryOne(ctx, &row, `FROM users WHERE type = $1 AND age < $2`, "fakeType", 42)
+			},
+			queryErr: errors.New("fakeErrMsg"),
+
+			expectLoggedQueryToContain: []string{"SELECT", "count", "type = $1"},
 			expectLoggedParams:         []interface{}{"fakeType", 42},
 			expectLoggedErrToContain:   []string{"fakeErrMsg"},
 		},
@@ -163,14 +206,25 @@ func TestInjectLogger(t *testing.T) {
 		t.Run(test.desc, func(t *testing.T) {
 			var inputQuery string
 			var inputParams []interface{}
+			numRows := 1
 			c := DB{
+				dialect: sqldialect.SupportedDialects["postgres"],
 				db: mockDBAdapter{
 					QueryContextFn: func(ctx context.Context, query string, params ...interface{}) (Rows, error) {
 						inputQuery = query
 						inputParams = params
 
 						return mockRows{
-							NextFn: func() bool { return false },
+							ScanFn: func(args ...interface{}) error {
+								return nil
+							},
+							// Make sure this mock will return a single row
+							// for the purposes of this test:
+							NextFn: func() bool {
+								numRows--
+								return numRows >= 0
+							},
+							ColumnsFn: func() ([]string, error) { return []string{"count"}, nil },
 						}, test.queryErr
 					},
 					ExecContextFn: func(ctx context.Context, query string, params ...interface{}) (Result, error) {
diff --git a/logger.go b/logger.go
index de3ef85..245c05c 100644
--- a/logger.go
+++ b/logger.go
@@ -1,6 +1,9 @@
 package ksql
 
-import "context"
+import (
+	"context"
+	"strings"
+)
 
 type loggerKey struct{}
 
@@ -15,7 +18,9 @@ func InjectLogger(
 	level string,
 	logFn func(ctx context.Context, values LogValues),
 ) context.Context {
+	level = strings.ToLower(level)
 	if level != "info" {
+		// Default to the least verbose level:
 		level = "error"
 	}