wait until StateHasSession

Author: shanth96

go/vt/topo/zk2topo/zk_conn.go

@@ -409,7 +409,7 @@ func dialZk(ctx context.Context, addr string) (*zk.Conn, <-chan zk.Event, error)
 			return nil, nil, ctx.Err()
 		case event := <-session:
 			switch event.State {
-			case zk.StateConnected:
+			case zk.StateHasSession:
 				// success
 				return zconn, session, nil
 

go/vt/topo/zk2topo/zk_conn_test.go

@@ -0,0 +1,26 @@
+package zk2topo
+
+import (
+	"context"
+	"testing"
+
+	"github.com/z-division/go-zookeeper/zk"
+	"vitess.io/vitess/go/testfiles"
+	"vitess.io/vitess/go/vt/zkctl"
+)
+
+func TestDialWaitsForSession(t *testing.T) {
+	zkd, serverAddr := zkctl.StartLocalZk(testfiles.GoVtTopoZk2topoZkID, testfiles.GoVtTopoZk2topoPort)
+	defer zkd.Teardown()
+
+	conn, _, err := dialZk(context.Background(), serverAddr)
+	if err != nil {
+		t.Fatalf("dialZk failed: %v", err)
+	}
+
+	defer conn.Close()
+
+	if conn.State() != zk.StateHasSession {
+		t.Fatalf("zk connection not in StateHasSession: %v", conn.State())
+	}
+}