From 9e8893a0b5325a76b2784958bbe743ff3e831401 Mon Sep 17 00:00:00 2001
From: Ruslan Drozhdzh <30860269+rdrozhdzh@users.noreply.github.com>
Date: Wed, 25 Apr 2018 21:15:49 +0300
Subject: [PATCH] rework TestProxyClose (#1735)

---
 plugin/forward/proxy_test.go | 30 +++++++++++++++++++++---------
 1 file changed, 21 insertions(+), 9 deletions(-)

diff --git a/plugin/forward/proxy_test.go b/plugin/forward/proxy_test.go
index a46b3f1ee..acd3d240c 100644
--- a/plugin/forward/proxy_test.go
+++ b/plugin/forward/proxy_test.go
@@ -2,8 +2,9 @@ package forward
 
 import (
 	"context"
-	"sync"
+	"runtime"
 	"testing"
+	"time"
 
 	"github.com/coredns/coredns/plugin/pkg/dnstest"
 	"github.com/coredns/coredns/plugin/test"
@@ -30,30 +31,41 @@ func TestProxyClose(t *testing.T) {
 		p := NewProxy(s.Addr, nil /* no TLS */)
 		p.start(hcDuration)
 
-		var wg sync.WaitGroup
-		wg.Add(5)
+		doneCnt := 0
+		doneCh := make(chan bool)
+		timeCh := time.After(10 * time.Second)
 		go func() {
 			p.connect(ctx, state, false, false)
-			wg.Done()
+			doneCh <- true
 		}()
 		go func() {
 			p.connect(ctx, state, true, false)
-			wg.Done()
+			doneCh <- true
 		}()
 		go func() {
 			p.close()
-			wg.Done()
+			doneCh <- true
 		}()
 		go func() {
 			p.connect(ctx, state, false, false)
-			wg.Done()
+			doneCh <- true
 		}()
 		go func() {
 			p.connect(ctx, state, true, false)
-			wg.Done()
+			doneCh <- true
 		}()
-		wg.Wait()
 
+		for doneCnt < 5 {
+			select {
+			case <-doneCh:
+				doneCnt++
+			case <-timeCh:
+				t.Error("TestProxyClose is running too long, dumping goroutines:")
+				buf := make([]byte, 100000)
+				stackSize := runtime.Stack(buf, true)
+				t.Fatal(string(buf[:stackSize]))
+			}
+		}
 		if p.inProgress != 0 {
 			t.Errorf("unexpected query in progress")
 		}