From 711c77311ecaccd9fced1b107191883866d9c6d1 Mon Sep 17 00:00:00 2001
From: Gyu-Ho Lee <gyuhox@gmail.com>
Date: Sat, 16 Jan 2016 13:27:21 -0800
Subject: [PATCH] integration: add TestV3WatchMultiple

For https://github.com/coreos/etcd/issues/4216.
---
 integration/v3_grpc_test.go | 88 +++++++++++++++++++++++++++++++++++++
 1 file changed, 88 insertions(+)

diff --git a/integration/v3_grpc_test.go b/integration/v3_grpc_test.go
index 7c96595ff698..2ff0d71fc247 100644
--- a/integration/v3_grpc_test.go
+++ b/integration/v3_grpc_test.go
@@ -491,3 +491,91 @@ func TestV3WatchCancel(t *testing.T) {
 
 	clus.Terminate(t)
 }
+
+// TestV3WatchMultiple tests multiple watchers on the same key
+// and one watcher with matching prefix. It first puts the key
+// that matches all watchers, and another key that matches only
+// one watcher to test if it receives expected events.
+func TestV3WatchMultiple(t *testing.T) {
+	clus := newClusterGRPC(t, &clusterConfig{size: 3})
+	wAPI := pb.NewWatchClient(clus.RandConn())
+	kvc := pb.NewKVClient(clus.RandConn())
+
+	wStream, errW := wAPI.Watch(context.TODO())
+	if errW != nil {
+		t.Fatalf("wAPI.Watch error: %v", errW)
+	}
+
+	keyN := 5
+	for i := 0; i < keyN-1; i++ {
+		if err := wStream.Send(&pb.WatchRequest{CreateRequest: &pb.WatchCreateRequest{Key: []byte("foo")}}); err != nil {
+			t.Fatalf("wStream.Send error: %v", err)
+		}
+	}
+	if err := wStream.Send(&pb.WatchRequest{CreateRequest: &pb.WatchCreateRequest{Prefix: []byte("fo")}}); err != nil {
+		t.Fatalf("wStream.Send error: %v", err)
+	}
+
+	if _, err := kvc.Put(context.TODO(), &pb.PutRequest{Key: []byte("foo"), Value: []byte("bar")}); err != nil {
+		t.Fatalf("couldn't put key (%v)", err)
+	}
+
+	keysReceived := []string{}
+	watchIds := make(map[int64]struct{})
+	for len(keysReceived) != keyN {
+		wresp, err := wStream.Recv()
+		if err != nil {
+			t.Fatalf("wStream.Recv error: %v", err)
+		}
+		if wresp.Created {
+			continue
+		}
+		watchIds[wresp.WatchId] = struct{}{}
+		for _, ev := range wresp.Events {
+			keysReceived = append(keysReceived, string(ev.Kv.Key))
+		}
+	}
+
+	if len(watchIds) != keyN {
+		t.Errorf("len(watchIds) got = %d, want = %d", len(watchIds), keyN)
+	}
+
+	keysToReceive := []string{"foo", "foo", "foo", "foo", "foo"}
+	if !reflect.DeepEqual(keysReceived, keysToReceive) {
+		t.Errorf("unexpected keys received %#q, want %#q", keysReceived, keysToReceive)
+	}
+
+	// now put one key that has only one matching watcher
+	if _, err := kvc.Put(context.TODO(), &pb.PutRequest{Key: []byte("fo"), Value: []byte("bar")}); err != nil {
+		t.Fatalf("couldn't put key (%v)", err)
+	}
+	wresp, err := wStream.Recv()
+	if err != nil {
+		t.Errorf("wStream.Recv error: %v", err)
+	}
+	if len(wresp.Events) != 1 {
+		t.Fatalf("len(wresp.Events) got = %d, want = 1", len(wresp.Events))
+	}
+	if string(wresp.Events[0].Kv.Key) != "fo" {
+		t.Errorf("wresp.Events[0].Kv.Key got = %s, want = fo", wresp.Events[0].Kv.Key)
+	}
+
+	// now Recv should block because there is no more events coming
+	rCh := make(chan *pb.WatchResponse)
+	go func() {
+		resp, _ := wStream.Recv()
+		rCh <- resp
+	}()
+	select {
+	case nr := <-rCh:
+		t.Errorf("unexpected response is received %+v", nr)
+	case <-time.After(2 * time.Second):
+	}
+	wStream.CloseSend()
+	rv, ok := <-rCh
+	if rv != nil || !ok {
+		t.Errorf("rv, ok got = %v %v, want = nil true", rv, ok)
+	}
+
+	clus.Terminate(t)
+}