From dc931aba7bcfff12bb56b2d5b749b2a223584fb7 Mon Sep 17 00:00:00 2001
From: Matt Lord <mattalord@gmail.com>
Date: Thu, 12 Dec 2024 20:24:52 -0500
Subject: [PATCH] Initial work on new unit test

Signed-off-by: Matt Lord <mattalord@gmail.com>
---
 .../vreplication/vplayer_flaky_test.go        | 92 ++++++++++++-------
 1 file changed, 59 insertions(+), 33 deletions(-)

diff --git a/go/vt/vttablet/tabletmanager/vreplication/vplayer_flaky_test.go b/go/vt/vttablet/tabletmanager/vreplication/vplayer_flaky_test.go
index 9c17780b66e..bfb6866c61a 100644
--- a/go/vt/vttablet/tabletmanager/vreplication/vplayer_flaky_test.go
+++ b/go/vt/vttablet/tabletmanager/vreplication/vplayer_flaky_test.go
@@ -1519,13 +1519,14 @@ func TestPlayerRowMove(t *testing.T) {
 	validateQueryCountStat(t, "replicate", 3)
 }
 
-/* TODO: build this out and get it working
 func TestPlayerUpdatePK(t *testing.T) {
 	defer deleteTablet(addTablet(100))
 
 	execStatements(t, []string{
-		"create table src(id int, bd blob, jd json, primary key(id))",
-		fmt.Sprintf("create table %s.dst(id int, bd blob, jd json, primary key(id))", vrepldb),
+		"set @@global.binlog_row_image=NOBLOB",
+		"set @@global.binlog_row_value_options=PARTIAL_JSON",
+		"create table src (id int, jd json, bd blob, primary key(id))",
+		fmt.Sprintf("create table %s.dst (id int, jd json, bd blob, primary key(id))", vrepldb),
 	})
 	defer execStatements(t, []string{
 		"drop table src",
@@ -1547,38 +1548,63 @@ func TestPlayerUpdatePK(t *testing.T) {
 	cancel, _ := startVReplication(t, bls, "")
 	defer cancel()
 
-	execStatements(t, []string{
-		"insert into src values(1, 'blob data', _utf8mb4'{\"key1\":\"val1\"}'), (2, 'blob data2', _utf8mb4'{\"key2\":\"val2\"}'), (3, 'blob data3', _utf8mb4'{\"key3\":\"val3\"}')",
-	})
-	expectDBClientQueries(t, qh.Expect(
-		"begin",
-		"insert into dst(id,bd,jd) values (1,_binary'blob data','{\"key1\": \"val1\"}'), (2,_binary'blob data2','{\"key2\": \"val2\"}'), (3,_binary'blob data3','{\"key3\": \"val3\"}')",
-		"/update _vt.vreplication set pos=",
-		"commit",
-	))
-	expectData(t, "dst", [][]string{
-		{"1", "1", "1"},
-		{"2", "5", "2"},
-	})
-	validateQueryCountStat(t, "replicate", 1)
+	testCases := []struct {
+		input  string
+		output string
+		data   [][]string
+	}{
+		{
+			input:  "insert into src (id, jd, bd) values (1, '{\"key1\": \"val1\"}', 'blob data'), (2, '{\"key2\": \"val2\"}', 'blob data2'), (3, '{\"key3\": \"val3\"}', 'blob data3')",
+			output: "insert into dst(id,jd,bd) values (1,'{\"key1\": \"val1\"}',_binary'blob data'), (2,'{\"key2\": \"val2\"}',_binary'blob data2'), (3,'{\"key3\": \"val3\"}',_binary'blob data3')",
+			data: [][]string{
+				{"1", "{\"key1\": \"val1\"}", "blob data"},
+				{"2", "{\"key2\": \"val2\"}", "blob data2"},
+				{"3", "{\"key3\": \"val3\"}", "blob data3"},
+			},
+		},
+		/*
+			{
+				input:  `update src set jd=JSON_SET(jd, '$.color', 'red') where id = 1`,
+				output: `update src set jd=JSON_SET(jd, '$.color', 'red') where id = 1`,
+				data: [][]string{
+					{"1", "{\"key1\": \"val1\", \"color\": \"red\"}", "blob data"},
+					{"2", "{\"key2\": \"val2\"}", "blob data2"},
+					{"3", "{\"key3\": \"val3\"}", "blob data3"},
+				},
+			},
+		*/
+	}
 
-	execStatements(t, []string{
-		"update src set val1=1, val2=4 where id=3",
-	})
-	expectDBClientQueries(t, qh.Expect(
-		"begin",
-		"update dst set sval2=sval2-ifnull(3, 0), rcount=rcount-1 where val1=2",
-		"insert into dst(val1,sval2,rcount) values (1,ifnull(4, 0),1) on duplicate key update sval2=sval2+ifnull(values(sval2), 0), rcount=rcount+1",
-		"/update _vt.vreplication set pos=",
-		"commit",
-	))
-	expectData(t, "dst", [][]string{
-		{"1", "5", "2"},
-		{"2", "2", "1"},
-	})
-	validateQueryCountStat(t, "replicate", 3)
+	for _, tc := range testCases {
+		execStatements(t, []string{tc.input})
+		want := qh.Expect(
+			"begin",
+			tc.output,
+			"/update _vt.vreplication set pos=",
+			"commit",
+		)
+		expectDBClientQueries(t, want)
+		expectData(t, "dst", tc.data)
+	}
+
+	/*
+		execStatements(t, []string{
+			"update src set val1=1, val2=4 where id=3",
+		})
+		expectDBClientQueries(t, qh.Expect(
+			"begin",
+			"update dst set sval2=sval2-ifnull(3, 0), rcount=rcount-1 where val1=2",
+			"insert into dst(val1,sval2,rcount) values (1,ifnull(4, 0),1) on duplicate key update sval2=sval2+ifnull(values(sval2), 0), rcount=rcount+1",
+			"/update _vt.vreplication set pos=",
+			"commit",
+		))
+		expectData(t, "dst", [][]string{
+			{"1", "5", "2"},
+			{"2", "2", "1"},
+		})
+		validateQueryCountStat(t, "replicate", 3)
+	*/
 }
-*/
 
 func TestPlayerTypes(t *testing.T) {
 	defer deleteTablet(addTablet(100))