From e1122eaf853de7e1748696802a105e689ca745fd Mon Sep 17 00:00:00 2001
From: Matt Lord <mattalord@gmail.com>
Date: Tue, 14 Jan 2025 14:18:06 -0500
Subject: [PATCH] Fix test when run with others in the package

Signed-off-by: Matt Lord <mattalord@gmail.com>
---
 .../endtoend/vtgate/vschema/vschema_test.go   | 37 ++++++++++++-------
 1 file changed, 23 insertions(+), 14 deletions(-)

diff --git a/go/test/endtoend/vtgate/vschema/vschema_test.go b/go/test/endtoend/vtgate/vschema/vschema_test.go
index 885697696aa..5f26061453f 100644
--- a/go/test/endtoend/vtgate/vschema/vschema_test.go
+++ b/go/test/endtoend/vtgate/vschema/vschema_test.go
@@ -25,6 +25,7 @@ import (
 	"path"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -168,10 +169,18 @@ func TestVSchema(t *testing.T) {
 	utils.AssertMatches(t, conn, "delete from vt_user", `[]`)
 
 	if utils.BinaryIsAtLeastAtVersion(22, "vtgate") {
+		// Don't allow any users to modify the vschema via the SQL API
+		// in order to test that behavior.
 		writeConfig(configFile, map[string]string{
 			"vschema_ddl_authorized_users": "",
 		})
-
+		// Allow anyone to modify the vschema via the SQL API again when
+		// the test completes.
+		defer func() {
+			writeConfig(configFile, map[string]string{
+				"vschema_ddl_authorized_users": "%",
+			})
+		}()
 		require.EventuallyWithT(t, func(t *assert.CollectT) {
 			_, err = conn.ExecuteFetch("ALTER VSCHEMA DROP TABLE main", 1000, false)
 			assert.Error(t, err)
@@ -183,13 +192,14 @@ func TestVSchema(t *testing.T) {
 // TestVSchemaSQLAPIConcurrency tests that we prevent lost writes when we have
 // concurrent vschema changes being made via the SQL API.
 func TestVSchemaSQLAPIConcurrency(t *testing.T) {
-	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second)
 	defer cancel()
 	conn, err := mysql.Connect(ctx, &vtParams)
 	require.NoError(t, err)
 	defer conn.Close()
 
-	utils.AssertMatches(t, conn, "SHOW VSCHEMA TABLES", `[]`)
+	initialVSchema, err := conn.ExecuteFetch("SHOW VSCHEMA TABLES", -1, false)
+	require.NoError(t, err)
 	baseTableName := "t"
 	numTables := 1000
 	mysqlConns := make([]*mysql.Conn, numTables)
@@ -201,26 +211,24 @@ func TestVSchemaSQLAPIConcurrency(t *testing.T) {
 	}
 
 	wg := sync.WaitGroup{}
-	preventedLostWrites := false
+	preventedLostWrites := atomic.Bool{}
 	for i := 0; i < numTables; i++ {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			time.Sleep(time.Duration(rand.Intn(1000) * int(time.Nanosecond)))
+			time.Sleep(time.Duration(rand.Intn(100) * int(time.Nanosecond)))
 			tableName := fmt.Sprintf("%s%d", baseTableName, i)
 			_, err = mysqlConns[i].ExecuteFetch(fmt.Sprintf("ALTER VSCHEMA ADD TABLE %s", tableName), -1, false)
-			if err != nil {
-				// The error we get is an SQL error so we have to do string matching.
-				if err != nil && strings.Contains(err.Error(), vtgate.ErrStaleVSchema.Error()) {
-					preventedLostWrites = true
-				}
+			// The error we get is an SQL error so we have to do string matching.
+			if err != nil && strings.Contains(err.Error(), vtgate.ErrStaleVSchema.Error()) {
+				preventedLostWrites.Store(true)
 			} else {
 				require.NoError(t, err)
-				time.Sleep(time.Duration(rand.Intn(1000) * int(time.Nanosecond)))
+				time.Sleep(time.Duration(rand.Intn(75) * int(time.Nanosecond)))
 				_, err = mysqlConns[i].ExecuteFetch(fmt.Sprintf("ALTER VSCHEMA DROP TABLE %s", tableName), -1, false)
 				// The error we get is an SQL error so we have to do string matching.
 				if err != nil && strings.Contains(err.Error(), vtgate.ErrStaleVSchema.Error()) {
-					preventedLostWrites = true
+					preventedLostWrites.Store(true)
 				} else {
 					require.NoError(t, err)
 				}
@@ -228,7 +236,7 @@ func TestVSchemaSQLAPIConcurrency(t *testing.T) {
 		}()
 	}
 	wg.Wait()
-	require.True(t, preventedLostWrites)
+	require.True(t, preventedLostWrites.Load())
 
 	// Cleanup any tables that were not dropped because the DROP query
 	// failed due to a bad node version.
@@ -236,5 +244,6 @@ func TestVSchemaSQLAPIConcurrency(t *testing.T) {
 		tableName := fmt.Sprintf("%s%d", baseTableName, i)
 		_, _ = mysqlConns[i].ExecuteFetch(fmt.Sprintf("ALTER VSCHEMA DROP TABLE %s", tableName), -1, false)
 	}
-	utils.AssertMatches(t, conn, "SHOW VSCHEMA TABLES", `[]`)
+	// Confirm that we're back to the initial state.
+	utils.AssertMatches(t, conn, "SHOW VSCHEMA TABLES", fmt.Sprintf("%v", initialVSchema.Rows))
 }