Added ability to wait for correct shutdown by waiting for closed channel

zk/cluster_test.go

@@ -2,6 +2,7 @@ package zk
 
 import (
 	"testing"
+	"time"
 )
 
 func TestBasicCluster(t *testing.T) {
@@ -55,3 +56,39 @@ func TestClientClusterFailover(t *testing.T) {
 		t.Fatal("Wrong data for node 2")
 	}
 }
+
+func TestWaitForClose(t *testing.T) {
+	ts, err := StartTestCluster(1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer ts.Stop()
+	zk, err := ts.Connect(0)
+	if err != nil {
+		t.Fatalf("Connect returned error: %+v", err)
+	}
+	timeout := time.After(30*time.Second)
+CONNECTED:
+	for{
+		select{
+		case ev := <-zk.eventChan :
+			if ev.State == StateConnected {
+				break CONNECTED;
+			}
+		case <-timeout:
+			zk.Close()
+			t.Fatal("Timeout")
+		}
+	}
+	zk.Close()
+	for{
+		select{
+		case _,ok := <-zk.eventChan :
+			if !ok {
+				return
+			}
+		case <-timeout:
+			t.Fatal("Timeout")
+		}
+	}
+}

zk/conn.go

@@ -142,6 +142,7 @@ func ConnectWithDialer(servers []string, recvTimeout time.Duration, dialer Diale
 		conn.loop()
 		conn.flushRequests(ErrClosing)
 		conn.invalidateWatches(ErrClosing)
+		close(conn.eventChan)
 	}()
 	return &conn, ec, nil
 }