From 6bd75223368bc23256d14b7a93480d2b88a8521f Mon Sep 17 00:00:00 2001
From: Gyu-Ho Lee <gyuhox@gmail.com>
Date: Fri, 14 Jul 2017 14:33:48 -0700
Subject: [PATCH] integration: test embed.Etcd.Close with watch

Ensure 'Close' returns in time when there are open
connections (watch streams).

Signed-off-by: Gyu-Ho Lee <gyuhox@gmail.com>
---
 integration/embed_test.go | 44 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 44 insertions(+)

diff --git a/integration/embed_test.go b/integration/embed_test.go
index 8cba0b39605c..c84852edf628 100644
--- a/integration/embed_test.go
+++ b/integration/embed_test.go
@@ -15,13 +15,16 @@
 package integration
 
 import (
+	"context"
 	"fmt"
 	"net/url"
 	"os"
 	"path/filepath"
 	"strings"
 	"testing"
+	"time"
 
+	"github.com/coreos/etcd/clientv3"
 	"github.com/coreos/etcd/embed"
 )
 
@@ -102,6 +105,47 @@ func TestEmbedEtcd(t *testing.T) {
 	}
 }
 
+// TestEmbedEtcdGracefulStop ensures embedded server stops
+// cutting existing transports.
+func TestEmbedEtcdGracefulStop(t *testing.T) {
+	cfg := embed.NewConfig()
+
+	urls := newEmbedURLs(2)
+	setupEmbedCfg(cfg, []url.URL{urls[0]}, []url.URL{urls[1]})
+
+	cfg.Dir = filepath.Join(os.TempDir(), fmt.Sprintf("embed-etcd"))
+	os.RemoveAll(cfg.Dir)
+	defer os.RemoveAll(cfg.Dir)
+
+	e, err := embed.StartEtcd(cfg)
+	if err != nil {
+		t.Fatal(err)
+	}
+	<-e.Server.ReadyNotify() // wait for e.Server to join the cluster
+
+	cli, err := clientv3.New(clientv3.Config{Endpoints: []string{urls[0].String()}})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer cli.Close()
+
+	// open watch connection
+	cli.Watch(context.Background(), "foo")
+
+	donec := make(chan struct{})
+	go func() {
+		e.Close()
+		close(donec)
+	}()
+	select {
+	case err := <-e.Err():
+		t.Fatal(err)
+	case <-donec:
+	case <-time.After(5 * time.Second):
+		t.Fatalf("took too long to close server")
+	}
+}
+
 func newEmbedURLs(n int) (urls []url.URL) {
 	for i := 0; i < n; i++ {
 		u, _ := url.Parse(fmt.Sprintf("unix://localhost:%d%06d", os.Getpid(), i))