From 92cd5293cf3668f7ae6d114d4592fbafd6894496 Mon Sep 17 00:00:00 2001
From: Tim Vaillancourt <tim@timvaillancourt.com>
Date: Thu, 20 Jun 2024 01:24:44 +0200
Subject: [PATCH] Add further testing to CountStar formatting

Signed-off-by: Tim Vaillancourt <tim@timvaillancourt.com>
---
 go/vt/sqlparser/ast_format.go           |  8 +++----
 go/vt/sqlparser/ast_format_fast.go      | 11 +++++-----
 go/vt/sqlparser/ast_format_fast_test.go | 28 +++++++++++++++++++++++++
 go/vt/sqlparser/ast_format_test.go      | 28 +++++++++++++++++++++++++
 4 files changed, 65 insertions(+), 10 deletions(-)
 create mode 100644 go/vt/sqlparser/ast_format_fast_test.go
 create mode 100644 go/vt/sqlparser/ast_format_test.go

diff --git a/go/vt/sqlparser/ast_format.go b/go/vt/sqlparser/ast_format.go
index 830420e610d..f2d35169419 100644
--- a/go/vt/sqlparser/ast_format.go
+++ b/go/vt/sqlparser/ast_format.go
@@ -2656,12 +2656,12 @@ func (node *Count) Format(buf *TrackedBuffer) {
 }
 
 func (node *CountStar) Format(buf *TrackedBuffer) {
+	name := node.AggrName()
 	if node.Name != "" {
-		buf.astPrintf(node, "%s(*)", node.Name)
-	} else {
-		buf.astPrintf(node, "%s(", node.AggrName())
-		buf.WriteString("*)")
+		name = node.Name
 	}
+	buf.astPrintf(node, "%s(", name)
+	buf.WriteString("*)")
 }
 
 func (node *Avg) Format(buf *TrackedBuffer) {
diff --git a/go/vt/sqlparser/ast_format_fast.go b/go/vt/sqlparser/ast_format_fast.go
index 8115288278e..636be927129 100644
--- a/go/vt/sqlparser/ast_format_fast.go
+++ b/go/vt/sqlparser/ast_format_fast.go
@@ -3482,14 +3482,13 @@ func (node *Count) formatFast(buf *TrackedBuffer) {
 }
 
 func (node *CountStar) formatFast(buf *TrackedBuffer) {
+	name := node.AggrName()
 	if node.Name != "" {
-		buf.WriteString(node.Name)
-		buf.WriteString("(*)")
-	} else {
-		buf.WriteString(node.AggrName())
-		buf.WriteByte('(')
-		buf.WriteString("*)")
+		name = node.Name
 	}
+	buf.WriteString(name)
+	buf.WriteByte('(')
+	buf.WriteString("*)")
 }
 
 func (node *Avg) formatFast(buf *TrackedBuffer) {
diff --git a/go/vt/sqlparser/ast_format_fast_test.go b/go/vt/sqlparser/ast_format_fast_test.go
new file mode 100644
index 00000000000..365ceb3f8b5
--- /dev/null
+++ b/go/vt/sqlparser/ast_format_fast_test.go
@@ -0,0 +1,28 @@
+package sqlparser
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCountStarFormatFast(t *testing.T) {
+	{
+		buf := NewTrackedBuffer(nil)
+		node := &CountStar{}
+		node.formatFast(buf)
+		assert.Equal(t, &ParsedQuery{Query: "count(*)"}, buf.ParsedQuery())
+	}
+	{
+		buf := NewTrackedBuffer(nil)
+		node := &CountStar{Name: "Count"}
+		node.formatFast(buf)
+		assert.Equal(t, &ParsedQuery{Query: "Count(*)"}, buf.ParsedQuery())
+	}
+	{
+		buf := NewTrackedBuffer(nil)
+		node := &CountStar{Name: "COUNT"}
+		node.formatFast(buf)
+		assert.Equal(t, &ParsedQuery{Query: "COUNT(*)"}, buf.ParsedQuery())
+	}
+}
diff --git a/go/vt/sqlparser/ast_format_test.go b/go/vt/sqlparser/ast_format_test.go
new file mode 100644
index 00000000000..e09ba7f4b5c
--- /dev/null
+++ b/go/vt/sqlparser/ast_format_test.go
@@ -0,0 +1,28 @@
+package sqlparser
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCountStarFormat(t *testing.T) {
+	{
+		buf := NewTrackedBuffer(nil)
+		node := &CountStar{}
+		node.Format(buf)
+		assert.Equal(t, &ParsedQuery{Query: "count(*)"}, buf.ParsedQuery())
+	}
+	{
+		buf := NewTrackedBuffer(nil)
+		node := &CountStar{Name: "Count"}
+		node.Format(buf)
+		assert.Equal(t, &ParsedQuery{Query: "Count(*)"}, buf.ParsedQuery())
+	}
+	{
+		buf := NewTrackedBuffer(nil)
+		node := &CountStar{Name: "COUNT"}
+		node.Format(buf)
+		assert.Equal(t, &ParsedQuery{Query: "COUNT(*)"}, buf.ParsedQuery())
+	}
+}