From 030ea5a2b16b01cf4cf994ef3ed9332e32d1c4be Mon Sep 17 00:00:00 2001
From: Matt Lord <mattalord@gmail.com>
Date: Sat, 11 Jan 2025 21:55:17 -0500
Subject: [PATCH] Improve unit test coverage

Signed-off-by: Matt Lord <mattalord@gmail.com>
---
 .../tabletmanager/vdiff/table_differ_test.go  | 127 ++++++++++++++++++
 1 file changed, 127 insertions(+)
 create mode 100644 go/vt/vttablet/tabletmanager/vdiff/table_differ_test.go

diff --git a/go/vt/vttablet/tabletmanager/vdiff/table_differ_test.go b/go/vt/vttablet/tabletmanager/vdiff/table_differ_test.go
new file mode 100644
index 00000000000..998d328c3c8
--- /dev/null
+++ b/go/vt/vttablet/tabletmanager/vdiff/table_differ_test.go
@@ -0,0 +1,127 @@
+/*
+Copyright 2025 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 vdiff
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"vitess.io/vitess/go/sqltypes"
+	"vitess.io/vitess/go/stats"
+	"vitess.io/vitess/go/vt/binlog/binlogplayer"
+
+	querypb "vitess.io/vitess/go/vt/proto/query"
+	tabletmanagerdatapb "vitess.io/vitess/go/vt/proto/tabletmanagerdata"
+)
+
+func TestUpdateTableProgress(t *testing.T) {
+	dr := &DiffReport{
+		TableName:     "test",
+		ProcessedRows: 1e9,
+	}
+	queryTemplate := `update _vt.vdiff_table set rows_compared = 1000000000, lastpk = '%s', report = '{"TableName":"test","ProcessedRows":1000000000,"MatchingRows":0,"MismatchedRows":0,"ExtraRowsSource":0,"ExtraRowsTarget":0}' where vdiff_id = 1 and table_name = 'test'`
+
+	testCases := []struct {
+		name           string
+		fields         []*querypb.Field
+		pkCols         []int
+		sourcePkCols   []int
+		lastRow        []sqltypes.Value
+		expectedLastPK string
+		wantErr        bool
+	}{
+		{
+			name: "identical PKs",
+			fields: []*querypb.Field{
+				{
+					Name: "a", Type: sqltypes.Int64,
+				},
+				{
+					Name: "b", Type: sqltypes.Int64,
+				},
+			},
+			pkCols:         []int{0, 1},
+			sourcePkCols:   []int{0, 1},
+			lastRow:        []sqltypes.Value{sqltypes.NewInt64(1), sqltypes.NewInt64(2)},
+			expectedLastPK: `target:{fields:{name:"a" type:INT64} fields:{name:"b" type:INT64} rows:{lengths:1 lengths:1 values:"12"}}`,
+		},
+		{
+			name: "more PK cols on target",
+			fields: []*querypb.Field{
+				{
+					Name: "a", Type: sqltypes.Int64,
+				},
+				{
+					Name: "b", Type: sqltypes.Int64,
+				},
+			},
+			pkCols:         []int{0, 1},
+			sourcePkCols:   []int{0},
+			lastRow:        []sqltypes.Value{sqltypes.NewInt64(1), sqltypes.NewInt64(2)},
+			expectedLastPK: `target:{fields:{name:"a" type:INT64} fields:{name:"b" type:INT64} rows:{lengths:1 lengths:1 values:"12"}} source:{fields:{name:"a" type:INT64} rows:{lengths:1 values:"1"}}`,
+		},
+		{
+			name: "more PK cols on source",
+			fields: []*querypb.Field{
+				{
+					Name: "a", Type: sqltypes.Int64,
+				},
+				{
+					Name: "b", Type: sqltypes.Int64,
+				},
+			},
+			pkCols:         []int{0},
+			sourcePkCols:   []int{0, 1},
+			lastRow:        []sqltypes.Value{sqltypes.NewInt64(1), sqltypes.NewInt64(2)},
+			expectedLastPK: `target:{fields:{name:"a" type:INT64} rows:{lengths:1 values:"1"}} source:{fields:{name:"a" type:INT64} fields:{name:"b" type:INT64} rows:{lengths:1 lengths:1 values:"12"}}`,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			dbc := binlogplayer.NewMockDBClient(t)
+			dbc.ExpectRequest(fmt.Sprintf(queryTemplate, tc.expectedLastPK), &sqltypes.Result{}, nil)
+			td := &tableDiffer{
+				wd: &workflowDiffer{
+					ct: &controller{
+						id:                 1,
+						TableDiffRowCounts: stats.NewCountersWithSingleLabel("", "", "Rows"),
+					},
+					opts: &tabletmanagerdatapb.VDiffOptions{
+						CoreOptions: &tabletmanagerdatapb.VDiffCoreOptions{},
+					},
+				},
+				table: &tabletmanagerdatapb.TableDefinition{
+					Name: "test",
+				},
+				tablePlan: &tablePlan{
+					pkCols:       tc.pkCols,
+					sourcePkCols: tc.sourcePkCols,
+					table: &tabletmanagerdatapb.TableDefinition{
+						Fields: tc.fields,
+					},
+				},
+			}
+			if err := td.updateTableProgress(dbc, dr, tc.lastRow); (err != nil) != tc.wantErr {
+				require.FailNow(t, "tableDiffer.updateTableProgress() error = %v, wantErr %v",
+					err, tc.wantErr)
+			}
+		})
+	}
+}