Merge pull request vitessio#5479 from planetscale/fix_e2e_test_fail

added fix for e2e intermittent test failure

go/test/endtoend/cluster/cluster_process.go

@@ -253,7 +253,13 @@ func (cluster *LocalProcessCluster) StartVtgate() (err error) {
 		cluster.VtGateExtraArgs)
 
 	log.Info(fmt.Sprintf("Vtgate started, connect to mysql using : mysql -h 127.0.0.1 -P %d", cluster.VtgateMySQLPort))
-	return cluster.VtgateProcess.Setup()
+	if err = cluster.VtgateProcess.Setup(); err != nil {
+		return err
+	}
+	if err = cluster.WaitForTabletsToHealthyInVtgate(); err != nil {
+		return err
+	}
+	return nil
 }
 
 // NewCluster instantiates a new cluster
@@ -282,6 +288,34 @@ func (cluster *LocalProcessCluster) ReStartVtgate() (err error) {
 	return err
 }
 
+// WaitForTabletsToHealthyInVtgate waits for all tablets in all shards to be healthy as per vtgate
+func (cluster *LocalProcessCluster) WaitForTabletsToHealthyInVtgate() (err error) {
+	var isRdOnlyPresent bool
+	for _, keyspace := range cluster.Keyspaces {
+		for _, shard := range keyspace.Shards {
+			isRdOnlyPresent = false
+			if err = cluster.VtgateProcess.WaitForStatusOfTabletInShard(fmt.Sprintf("%s.%s.master", keyspace.Name, shard.Name)); err != nil {
+				return err
+			}
+			if err = cluster.VtgateProcess.WaitForStatusOfTabletInShard(fmt.Sprintf("%s.%s.replica", keyspace.Name, shard.Name)); err != nil {
+				return err
+			}
+			for _, tablet := range shard.Vttablets {
+				if tablet.Type == "rdonly" {
+					isRdOnlyPresent = true
+				}
+			}
+			if isRdOnlyPresent {
+				err = cluster.VtgateProcess.WaitForStatusOfTabletInShard(fmt.Sprintf("%s.%s.rdonly", keyspace.Name, shard.Name))
+			}
+			if err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
 // Teardown brings down the cluster by invoking teardown for individual processes
 func (cluster *LocalProcessCluster) Teardown() (err error) {
 	if err = cluster.VtgateProcess.TearDown(); err != nil {

go/test/endtoend/cluster/vtgate_process.go

@@ -119,6 +119,18 @@ func (vtgate *VtgateProcess) Setup() (err error) {
 
 // WaitForStatus function checks if vtgate process is up and running
 func (vtgate *VtgateProcess) WaitForStatus() bool {
+	resp, err := http.Get(vtgate.VerifyURL)
+	if err != nil {
+		return false
+	}
+	if resp.StatusCode == 200 {
+		return true
+	}
+	return false
+}
+
+// GetStatusForTabletOfShard function gets status for a specific tablet of a shard in keyspace
+func (vtgate *VtgateProcess) GetStatusForTabletOfShard(name string) bool {
 	resp, err := http.Get(vtgate.VerifyURL)
 	if err != nil {
 		return false
@@ -134,17 +146,35 @@ func (vtgate *VtgateProcess) WaitForStatus() bool {
 		masterConnectionExist := false
 		if object.Kind() == reflect.Map {
 			for _, key := range object.MapKeys() {
-
-				if strings.Contains(key.String(),"master") {
-					masterConnectionExist = true
+				if key.String() == name {
+					value := fmt.Sprintf("%v", object.MapIndex(key))
+					return value == "1"
 				}
+
 			}
 		}
 		return masterConnectionExist
 	}
 	return false
 }
 
+// WaitForStatusOfTabletInShard function waits till status of a tablet in shard is 1
+func (vtgate *VtgateProcess) WaitForStatusOfTabletInShard(name string) error {
+	timeout := time.Now().Add(10 * time.Second)
+	for time.Now().Before(timeout) {
+		if vtgate.GetStatusForTabletOfShard(name) {
+			return nil
+		}
+		select {
+		case err := <-vtgate.exit:
+			return fmt.Errorf("process '%s' exited prematurely (err: %s)", vtgate.Name, err)
+		default:
+			time.Sleep(300 * time.Millisecond)
+		}
+	}
+	return fmt.Errorf("wait for %s failed", name)
+}
+
 // TearDown shuts down the running vtgate service
 func (vtgate *VtgateProcess) TearDown() error {
 	if vtgate.proc == nil || vtgate.exit == nil {

go/test/endtoend/clustertest/add_keyspace_test.go

@@ -64,6 +64,7 @@ func TestAddKeyspace(t *testing.T) {
 	// Restart vtgate process
 	_ = clusterInstance.VtgateProcess.TearDown()
 	_ = clusterInstance.VtgateProcess.Setup()
+	clusterInstance.WaitForTabletsToHealthyInVtgate()
 
 	ctx := context.Background()
 	vtParams := mysql.ConnParams{