update test to be more complete

Signed-off-by: Nick Van Wiggeren <[email protected]>

Author: nickvanw

go/test/endtoend/vtgate/misc_test.go

@@ -847,59 +847,63 @@ func TestTabletTargeting(t *testing.T) {
 	qr := utils.Exec(t, conn, "show vitess_tablets")
 	require.Greater(t, len(qr.Rows), 0, "no tablets found")
 
-	// Find PRIMARY tablets for both shards (-80 and 80-)
-	var primaryShard80Minus string // We'll target this shard
-	var primaryShard80Plus string  // We'll verify this shard gets no writes
+	// Find PRIMARY and REPLICA tablets for both shards (-80 and 80-)
+	var primaryShard80Minus string
+	var primaryShard80Plus string
 	for _, row := range qr.Rows {
 		shard := row[2].ToString()
 		tabletType := row[3].ToString()
-		if tabletType == "PRIMARY" {
-			switch shard {
-			case "-80":
-				primaryShard80Minus = row[0].ToString()
-			case "80-":
-				primaryShard80Plus = row[0].ToString()
-			}
+		switch {
+		case tabletType == "PRIMARY" && shard == "-80":
+			primaryShard80Minus = row[0].ToString()
+		case tabletType == "PRIMARY" && shard == "80-":
+			primaryShard80Plus = row[0].ToString()
 		}
 	}
 	require.NotEmpty(t, primaryShard80Minus, "no PRIMARY tablet found for -80 shard")
 	require.NotEmpty(t, primaryShard80Plus, "no PRIMARY tablet found for 80- shard")
 
-	// Test: Target a specific tablet in -80 shard and insert rows
-	// id1=1 and id1=2 both hash to -80 shard - this validates normal behavior
-	utils.Exec(t, conn, fmt.Sprintf("USE ks:-80@%s", primaryShard80Minus))
+	// Test 1: Shard targeting bypasses vindex resolution
+	// Insert data that would normally hash to 80- shard, but goes to -80 because of shard targeting
+	utils.Exec(t, conn, fmt.Sprintf("USE `ks:-80@%s`", primaryShard80Minus))
 	utils.Exec(t, conn, "insert into t1(id1, id2) values(1, 100), (2, 200)")
-	utils.AssertMatches(t, conn, "select id1 from t1 where id1 in (1, 2) order by id1", "[[INT64(1)] [INT64(2)]]")
-
-	// Test: Insert data that would normally hash to 80- shard, but goes to -80 because of shard targeting
 	// id1=4 hashes to 80-, but we're targeting -80 shard explicitly
-	utils.Exec(t, conn, fmt.Sprintf("USE ks:-80@%s", primaryShard80Minus))
 	utils.Exec(t, conn, "insert into t1(id1, id2) values(4, 400)")
 
 	// Verify the data went to -80 shard (not where vindex would have put it)
-	utils.Exec(t, conn, "USE ks:-80")
-	utils.AssertMatches(t, conn, "select id1 from t1 where id1=4", "[[INT64(4)]]")
+	utils.Exec(t, conn, "USE `ks:-80`")
+	utils.AssertMatches(t, conn, "select id1 from t1 where id1 in (1, 2, 4) order by id1", "[[INT64(1)] [INT64(2)] [INT64(4)]]")
 
-	// Verify the data did NOT go to 80- shard (where vindex says it should be)
-	utils.Exec(t, conn, "USE ks:80-")
+	// Verify the data did NOT go to 80- shard (where vindex says id1=4 should be)
+	utils.Exec(t, conn, "USE `ks:80-`")
 	utils.AssertIsEmpty(t, conn, "select id1 from t1 where id1=4")
 
-	// Test: Transaction with tablet-specific routing
-	utils.Exec(t, conn, fmt.Sprintf("USE ks:-80@%s", primaryShard80Minus))
+	// Test 2: Transaction with tablet-specific routing maintains sticky connection
+	utils.Exec(t, conn, fmt.Sprintf("USE `ks:-80@%s`", primaryShard80Minus))
 	utils.Exec(t, conn, "begin")
 	utils.Exec(t, conn, "insert into t1(id1, id2) values(10, 300)")
+	// Subsequent queries in transaction should go to same tablet
+	utils.AssertMatches(t, conn, "select id1 from t1 where id1=10", "[[INT64(10)]]")
 	utils.Exec(t, conn, "commit")
 	utils.AssertMatches(t, conn, "select id1 from t1 where id1=10", "[[INT64(10)]]")
 
-	// Test: Rollback with tablet-specific routing
-	utils.Exec(t, conn, fmt.Sprintf("USE ks:-80@%s", primaryShard80Minus))
+	// Test 5: Rollback with tablet-specific routing
+	utils.Exec(t, conn, fmt.Sprintf("USE `ks:-80@%s`", primaryShard80Minus))
 	utils.Exec(t, conn, "begin")
 	utils.Exec(t, conn, "insert into t1(id1, id2) values(20, 500)")
 	utils.Exec(t, conn, "rollback")
-	// 20 should not exist
 	utils.AssertIsEmpty(t, conn, "select id1 from t1 where id1=20")
 
-	// Test: Clear tablet targeting returns to normal routing
+	// Test 3: Invalid tablet alias should fail
+	_, err = conn.ExecuteFetch("USE `ks:-80@nonexistent-tablet`", 1, false)
+	require.Error(t, err, "invalid tablet alias should fail")
+
+	// Test 7: Tablet targeting without shard specification
+	utils.Exec(t, conn, fmt.Sprintf("USE `ks@%s`", primaryShard80Minus))
+	// Should be able to query data from the tablet's shard
+	utils.AssertMatches(t, conn, "select id1 from t1 where id1 in (1, 2, 4, 10) order by id1", "[[INT64(1)] [INT64(2)] [INT64(4)] [INT64(10)]]")
+
+	// Test 5: Clear tablet targeting returns to normal routing
 	utils.Exec(t, conn, "USE ks")
 	utils.AssertMatches(t, conn, "select id1 from t1 where id1 in (1, 2, 4, 10) order by id1", "[[INT64(1)] [INT64(2)] [INT64(4)] [INT64(10)]]")