From 3c8bc33431dce86fe1555567774f5a0adfef6c04 Mon Sep 17 00:00:00 2001
From: Manan Gupta <35839558+GuptaManan100@users.noreply.github.com>
Date: Tue, 3 Oct 2023 18:22:24 +0530
Subject: [PATCH] FlakyFix: `TestZk2Topo` (#14162)

Signed-off-by: Manan Gupta <manan@planetscale.com>
---
 go/vt/topo/zk2topo/server_test.go | 30 ++++++++++++++++++++++--------
 1 file changed, 22 insertions(+), 8 deletions(-)

diff --git a/go/vt/topo/zk2topo/server_test.go b/go/vt/topo/zk2topo/server_test.go
index a35e84f2950..9fd36740913 100644
--- a/go/vt/topo/zk2topo/server_test.go
+++ b/go/vt/topo/zk2topo/server_test.go
@@ -21,6 +21,9 @@ import (
 	"fmt"
 	"path"
 	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
 
 	"vitess.io/vitess/go/testfiles"
 	"vitess.io/vitess/go/vt/topo"
@@ -51,14 +54,25 @@ func TestZk2Topo(t *testing.T) {
 		// Note we exercise the observer feature here by passing in
 		// the same server twice, with a "|" separator.
 		ts, err := topo.OpenServer("zk2", serverAddr+"|"+serverAddr, globalRoot)
-		if err != nil {
-			t.Fatalf("OpenServer() failed: %v", err)
-		}
-		if err := ts.CreateCellInfo(context.Background(), test.LocalCellName, &topodatapb.CellInfo{
-			ServerAddress: serverAddr,
-			Root:          cellRoot,
-		}); err != nil {
-			t.Fatalf("CreateCellInfo() failed: %v", err)
+		require.NoError(t, err, "OpenServer() failed")
+		// We retry creating the cell info until we no longer get a connection error.
+		timeout := time.After(15 * time.Second)
+		for {
+			err = ts.CreateCellInfo(context.Background(), test.LocalCellName, &topodatapb.CellInfo{
+				ServerAddress: serverAddr,
+				Root:          cellRoot,
+			})
+			if err == nil {
+				break
+			}
+			select {
+			case <-timeout:
+				t.Fatalf("Timedout creating cell info - %v", err)
+				return nil
+			default:
+				require.ErrorContainsf(t, err, "could not connect to a server", "Received an error that isn't a connection error")
+				time.Sleep(1 * time.Second)
+			}
 		}
 
 		return ts