Add testing for shard scoped foreign keys (#15571)

Signed-off-by: Manan Gupta <manan@planetscale.com>

go/test/endtoend/vtgate/foreignkey/fk_fuzz_test.go

@@ -257,20 +257,20 @@ func (fz *fuzzer) generateDeleteDMLQuery() string {
 }
 
 // start starts running the fuzzer.
-func (fz *fuzzer) start(t *testing.T, sharded bool) {
+func (fz *fuzzer) start(t *testing.T, keyspace string) {
  // We mark the fuzzer thread to be running now.
  fz.shouldStop.Store(false)
  fz.wg.Add(fz.concurrency)
  for i := 0; i < fz.concurrency; i++ {
  fuzzerThreadId := i
  go func() {
- fz.runFuzzerThread(t, sharded, fuzzerThreadId)
+ fz.runFuzzerThread(t, keyspace, fuzzerThreadId)
  }()
  }
 }
 
 // runFuzzerThread is used to run a thread of the fuzzer.
-func (fz *fuzzer) runFuzzerThread(t *testing.T, sharded bool, fuzzerThreadId int) {
+func (fz *fuzzer) runFuzzerThread(t *testing.T, keyspace string, fuzzerThreadId int) {
  // Whenever we finish running this thread, we should mark the thread has stopped.
  defer func() {
  fz.wg.Done()
@@ -293,16 +293,9 @@ func (fz *fuzzer) runFuzzerThread(t *testing.T, sharded bool, fuzzerThreadId int
  defer mysqlDb.Close()
  }
  // Set the correct keyspace to use from VtGates.
- if sharded {
- _ = utils.Exec(t, mcmp.VtConn, "use `ks`")
- if vitessDb != nil {
- _, _ = vitessDb.Exec("use `ks`")
- }
- } else {
- _ = utils.Exec(t, mcmp.VtConn, "use `uks`")
- if vitessDb != nil {
- _, _ = vitessDb.Exec("use `uks`")
- }
+ _ = utils.Exec(t, mcmp.VtConn, fmt.Sprintf("use `%v`", keyspace))
+ if vitessDb != nil {
+ _, _ = vitessDb.Exec(fmt.Sprintf("use `%v`", keyspace))
  }
  if fz.queryFormat == OlapSQLQueries {
  _ = utils.Exec(t, mcmp.VtConn, "set workload = olap")
@@ -715,21 +708,19 @@ func TestFkFuzzTest(t *testing.T) {
  valFalse := false
  for _, fkState := range []*bool{nil, &valTrue, &valFalse} {
  for _, tt := range testcases {
- for _, testSharded := range []bool{false, true} {
+ for _, keyspace := range []string{unshardedKs, shardedKs} {
  for _, queryFormat := range []QueryFormat{OlapSQLQueries, SQLQueries, PreparedStatmentQueries, PreparedStatementPacket} {
  if fkState != nil && (queryFormat != SQLQueries || tt.concurrency != 1) {
  continue
  }
- t.Run(getTestName(tt.name, testSharded)+fmt.Sprintf(" FkState - %v QueryFormat - %v", sqlparser.FkChecksStateString(fkState), queryFormat), func(t *testing.T) {
+ t.Run(getTestName(tt.name, keyspace)+fmt.Sprintf(" FkState - %v QueryFormat - %v", sqlparser.FkChecksStateString(fkState), queryFormat), func(t *testing.T) {
  mcmp, closer := start(t)
  defer closer()
- // Set the correct keyspace to use from VtGates.
- if testSharded {
+ if keyspace == shardedKs {
  t.Skip("Skip test since we don't have sharded foreign key support yet")
- _ = utils.Exec(t, mcmp.VtConn, "use `ks`")
- } else {
- _ = utils.Exec(t, mcmp.VtConn, "use `uks`")
  }
+ // Set the correct keyspace to use from VtGates.
+ _ = utils.Exec(t, mcmp.VtConn, fmt.Sprintf("use `%v`", keyspace))
 
  // Ensure that the Vitess database is originally empty
  ensureDatabaseState(t, mcmp.VtConn, true)
@@ -739,7 +730,7 @@ func TestFkFuzzTest(t *testing.T) {
  fz := newFuzzer(tt.concurrency, tt.maxValForId, tt.maxValForCol, tt.insertShare, tt.deleteShare, tt.updateShare, queryFormat, fkState)
 
  // Start the fuzzer.
- fz.start(t, testSharded)
+ fz.start(t, keyspace)
 
  // Wait for the timeForTesting so that the threads continue to run.
  totalTime := time.After(tt.timeForTesting)