make vtbenchstat testable and write unit_tests

Signed-off-by: Andres Taylor <[email protected]>

src/cmd/vtbenchstat/main.go

@@ -37,9 +37,8 @@ func main() {
 	}
 
 	if len(traces) == 1 {
-		printSummary(traces[0])
+		printSummary(os.Stdout, terminalWidth(), highlightQuery, traces[0])
 	} else {
-		compareTraces(traces[0], traces[1])
+		compareTraces(os.Stdout, terminalWidth(), highlightQuery, traces[0], traces[1])
 	}
 }
-

src/cmd/vtbenchstat/vtbenchstat.go

@@ -22,6 +22,7 @@ import (
 	"github.com/alecthomas/chroma/quick"
 	"github.com/olekukonko/tablewriter"
 	"golang.org/x/term"
+	"io"
 	"math"
 	"os"
 	"sort"
@@ -180,16 +181,15 @@ func limitQueryLength(query string, termWidth int) string {
 	return processedQuery
 }
 
-func printSummary(file TraceFile) {
+func printSummary(out io.Writer, termWidth int, highLighter Highlighter, file TraceFile) {
 	summary := summarizeTraces(file)
-	termWidth, _, err := term.GetSize(int(os.Stdout.Fd()))
-	if err != nil {
-		termWidth = 80 // default to 80 if we can't get the terminal width
-	}
-	for _, query := range file.Queries {
+	for i, query := range file.Queries {
+		if i > 0 {
+			fmt.Fprintln(out)
+		}
 		querySummary := summary[query.Query]
-		printQuery(query, termWidth, false)
-		table := tablewriter.NewWriter(os.Stdout)
+		printQuery(out, termWidth, highLighter, query, false)
+		table := tablewriter.NewWriter(out)
 		table.SetAutoFormatHeaders(false)
 		table.SetHeader([]string{
 			"Route Calls",
@@ -204,33 +204,46 @@ func printSummary(file TraceFile) {
 			strconv.Itoa(querySummary.ShardsQueried),
 		})
 		table.Render()
-		fmt.Println()
 	}
 }
 
-func printQuery(q TracedQuery, terminalWidth int, significant bool) {
-	fmt.Printf("%s", queryPrefix)
-	err := quick.Highlight(os.Stdout, limitQueryLength(q.Query, terminalWidth), "sql", "terminal", "monokai")
+type Highlighter func(out io.Writer, query string) error
+
+func highlightQuery(out io.Writer, query string) error {
+	return quick.Highlight(out, query, "sql", "terminal", "monokai")
+}
+
+func noHighlight(out io.Writer, query string) error {
+	_, err := fmt.Fprint(out, query)
+	return err
+}
+
+func printQuery(out io.Writer, terminalWidth int, highLighter Highlighter, q TracedQuery, significant bool) {
+	fmt.Fprintf(out, "%s", queryPrefix)
+	err := highLighter(out, limitQueryLength(q.Query, terminalWidth))
 	if err != nil {
 		return
 	}
 	improved := ""
 	if significant {
 		improved = " (significant)"
 	}
-	fmt.Printf("\nLine # %s%s\n", q.LineNumber, improved)
+	fmt.Fprintf(out, "\nLine # %s%s\n", q.LineNumber, improved)
 }
 
 const significantChangeThreshold = 10
 
-func compareTraces(file1, file2 TraceFile) {
-	summary1 := summarizeTraces(file1)
-	summary2 := summarizeTraces(file2)
-
+func terminalWidth() int {
 	termWidth, _, err := term.GetSize(int(os.Stdout.Fd()))
 	if err != nil {
-		termWidth = 80 // default to 80 if we can't get the terminal width
+		return 80 // default to 80 if we can't get the terminal width
 	}
+	return termWidth
+}
+
+func compareTraces(out io.Writer, termWidth int, highLighter Highlighter, file1, file2 TraceFile) {
+	summary1 := summarizeTraces(file1)
+	summary2 := summarizeTraces(file2)
 
 	var significantChanges, totalQueries int
 	var s1RouteCalls, s1DataSent, s1MemoryRows, s1ShardsQueried int
@@ -244,7 +257,7 @@ func compareTraces(file1, file2 TraceFile) {
 		}
 		totalQueries++
 
-		table := tablewriter.NewWriter(os.Stdout)
+		table := tablewriter.NewWriter(out)
 		table.SetHeader([]string{"Metric", file1.Name, file2.Name, "Diff", "% Change"})
 		table.SetAutoFormatHeaders(false)
 
@@ -268,9 +281,9 @@ func compareTraces(file1, file2 TraceFile) {
 		s2MemoryRows += s2.RowsInMemory
 		s2ShardsQueried += s2.ShardsQueried
 
-		printQuery(s1.Q, termWidth, significant)
+		printQuery(out, termWidth, highLighter, s1.Q, significant)
 		table.Render()
-		fmt.Println()
+		fmt.Fprintln(out)
 	}
 
 	totalRouteCallsChange := float64(s2RouteCalls-s1RouteCalls) / float64(s1RouteCalls) * 100
@@ -279,12 +292,12 @@ func compareTraces(file1, file2 TraceFile) {
 	totalShardsQueriedChange := float64(s2ShardsQueried-s1ShardsQueried) / float64(s1ShardsQueried) * 100
 
 	// Print summary
-	fmt.Println("Summary:")
-	fmt.Printf("- %d out of %d queries showed significant change\n", significantChanges, totalQueries)
-	fmt.Printf("- Average change in Route Calls: %.2f%%\n", totalRouteCallsChange)
-	fmt.Printf("- Average change in Data Sent: %.2f%%\n", totalDataSentChange)
-	fmt.Printf("- Average change in Rows In Memory: %.2f%%\n", totalMemoryRowsChange)
-	fmt.Printf("- Average change in Shards Queried: %.2f%%\n", totalShardsQueriedChange)
+	fmt.Fprintln(out, "Summary:")
+	fmt.Fprintf(out, "- %d out of %d queries showed significant change\n", significantChanges, totalQueries)
+	fmt.Fprintf(out, "- Average change in Route Calls: %.2f%%\n", totalRouteCallsChange)
+	fmt.Fprintf(out, "- Average change in Data Sent: %.2f%%\n", totalDataSentChange)
+	fmt.Fprintf(out, "- Average change in Rows In Memory: %.2f%%\n", totalMemoryRowsChange)
+	fmt.Fprintf(out, "- Average change in Shards Queried: %.2f%%\n", totalShardsQueriedChange)
 }
 
 // compareMetric compares two metrics and appends the result to the table, returning true if the change is significant

src/cmd/vtbenchstat/vtbenchstat_test.go

@@ -0,0 +1,180 @@
+/*
+Copyright 2024 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package main
+
+import (
+	"github.com/stretchr/testify/assert"
+	"strings"
+	"testing"
+)
+
+var tf1 = TraceFile{
+	Name: "test",
+	Queries: []TracedQuery{{
+		Query:      "select * from music",
+		LineNumber: "1",
+		Trace: Trace{
+			OperatorType:       "Route",
+			Variant:            "Scatter",
+			NoOfCalls:          1,
+			AvgNumberOfRows:    16,
+			MedianNumberOfRows: 16,
+			ShardsQueried:      8,
+		},
+	}, {
+		Query:      "select tbl.foo, tbl2.bar from tbl join tbl2 on tbl.id = tbl2.id order by tbl.baz",
+		LineNumber: "2",
+		Trace: Trace{
+			OperatorType:       "Sort",
+			Variant:            "Memory",
+			NoOfCalls:          1,
+			AvgNumberOfRows:    16,
+			MedianNumberOfRows: 16,
+			Inputs: []Trace{{
+				OperatorType:       "Join",
+				Variant:            "Apply",
+				NoOfCalls:          1,
+				AvgNumberOfRows:    16,
+				MedianNumberOfRows: 16,
+				Inputs: []Trace{{
+					OperatorType:       "Route",
+					Variant:            "Scatter",
+					NoOfCalls:          1,
+					AvgNumberOfRows:    10,
+					MedianNumberOfRows: 10,
+					ShardsQueried:      8,
+				}, {
+					OperatorType:       "Route",
+					Variant:            "EqualUnique",
+					NoOfCalls:          10,
+					AvgNumberOfRows:    1,
+					MedianNumberOfRows: 1,
+					ShardsQueried:      10,
+				}},
+			}},
+		},
+	}},
+}
+var tf2 = TraceFile{
+	Name: "test",
+	Queries: []TracedQuery{{
+		Query:      "select * from music",
+		LineNumber: "1",
+		Trace: Trace{
+			OperatorType:       "Route",
+			Variant:            "Scatter",
+			NoOfCalls:          1,
+			AvgNumberOfRows:    16,
+			MedianNumberOfRows: 16,
+			ShardsQueried:      7,
+		},
+	}, {
+		Query:      "select tbl.foo, tbl2.bar from tbl join tbl2 on tbl.id = tbl2.id order by tbl.baz",
+		LineNumber: "2",
+		Trace: Trace{
+			OperatorType:       "Route",
+			Variant:            "Scatter",
+			NoOfCalls:          1,
+			AvgNumberOfRows:    16,
+			MedianNumberOfRows: 16,
+			ShardsQueried:      8,
+		},
+	}},
+}
+
+func TestSummary(t *testing.T) {
+	t.Run("tf1", func(t *testing.T) {
+		sb := &strings.Builder{}
+		printSummary(sb, 80, noHighlight, tf1)
+		assert.Equal(t, `Query: select * from music
+Line # 1
++-------------+-----------+----------------+----------------+
+| Route Calls | Rows Sent | Rows In Memory | Shards Queried |
++-------------+-----------+----------------+----------------+
+|           1 |        16 |              0 |              8 |
++-------------+-----------+----------------+----------------+
+
+Query: select tbl.foo, tbl2.bar from tbl join tbl2 on tbl.id = tbl2.id order ...
+Line # 2
++-------------+-----------+----------------+----------------+
+| Route Calls | Rows Sent | Rows In Memory | Shards Queried |
++-------------+-----------+----------------+----------------+
+|          11 |        20 |             16 |             18 |
++-------------+-----------+----------------+----------------+
+`, sb.String())
+	})
+
+	t.Run("tf2", func(t *testing.T) {
+		sb := &strings.Builder{}
+		printSummary(sb, 80, noHighlight, tf2)
+		assert.Equal(t, `Query: select * from music
+Line # 1
++-------------+-----------+----------------+----------------+
+| Route Calls | Rows Sent | Rows In Memory | Shards Queried |
++-------------+-----------+----------------+----------------+
+|           1 |        16 |              0 |              7 |
++-------------+-----------+----------------+----------------+
+
+Query: select tbl.foo, tbl2.bar from tbl join tbl2 on tbl.id = tbl2.id order ...
+Line # 2
++-------------+-----------+----------------+----------------+
+| Route Calls | Rows Sent | Rows In Memory | Shards Queried |
++-------------+-----------+----------------+----------------+
+|           1 |        16 |              0 |              8 |
++-------------+-----------+----------------+----------------+
+`, sb.String())
+	})
+}
+
+func TestCompareFiles(t *testing.T) {
+	sb := &strings.Builder{}
+	compareTraces(sb, 80, noHighlight, tf1, tf2)
+	s := sb.String()
+	want := `Query: select * from music
+Line # 1 (significant)
++----------------+------+------+------+----------+
+|     Metric     | test | test | Diff | % Change |
++----------------+------+------+------+----------+
+| Route Calls    |    1 |    1 |    0 | 0.00%    |
+| Rows Sent      |   16 |   16 |    0 | 0.00%    |
+| Rows In Memory |    0 |    0 |    0 | NaN%     |
+| Shards Queried |    8 |    7 |   -1 | -12.50%  |
++----------------+------+------+------+----------+
+
+Query: select tbl.foo, tbl2.bar from tbl join tbl2 on tbl.id = tbl2.id order ...
+Line # 2 (significant)
++----------------+------+------+------+----------+
+|     Metric     | test | test | Diff | % Change |
++----------------+------+------+------+----------+
+| Route Calls    |   11 |    1 |  -10 | -90.91%  |
+| Rows Sent      |   20 |   16 |   -4 | -20.00%  |
+| Rows In Memory |   16 |    0 |  -16 | -100.00% |
+| Shards Queried |   18 |    8 |  -10 | -55.56%  |
++----------------+------+------+------+----------+
+
+Summary:
+- 2 out of 2 queries showed significant change
+- Average change in Route Calls: -83.33%
+- Average change in Data Sent: -11.11%
+- Average change in Rows In Memory: -100.00%
+- Average change in Shards Queried: -42.31%
+`
+	if s != want {
+		panic("oh noes")
+	}
+	assert.Equal(t, want, s)
+}