wait until StateHasSession

Author: shanth96

go/vt/topo/zk2topo/zk_conn.go

@@ -406,7 +406,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

@@ -60,3 +60,19 @@ func TestZkConnClosedOnDisconnect(t *testing.T) {
 
 	require.Equal(t, zk.StateDisconnected, oldConn.State(), "Connection is not in disconnected state")
 }
+
+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())
+	}
+}