From 65b2aa0a631c007238ee23b56d25e88da19815d5 Mon Sep 17 00:00:00 2001
From: Tom Thorogood <me+github@tomthorogood.co.uk>
Date: Mon, 3 Dec 2018 18:10:52 +1030
Subject: [PATCH] Make TestTimeout less flaky (#863)

This test currently fails on occasion due to what appears to be
goroutine scheduling differences. This should eliminate those failures.
---
 client_test.go | 57 +++++++++++++++++++-------------------------------
 1 file changed, 21 insertions(+), 36 deletions(-)

diff --git a/client_test.go b/client_test.go
index b88c6f5c..311c5cea 100644
--- a/client_test.go
+++ b/client_test.go
@@ -485,49 +485,34 @@ func TestTimeout(t *testing.T) {
 	m := new(Msg)
 	m.SetQuestion("miek.nl.", TypeTXT)
 
-	// Use a channel + timeout to ensure we don't get stuck if the
-	// Client Timeout is not working properly
-	done := make(chan struct{}, 2)
+	runTest := func(name string, exchange func(m *Msg, addr string, timeout time.Duration) (*Msg, time.Duration, error)) {
+		t.Run(name, func(t *testing.T) {
+			start := time.Now()
 
-	timeout := time.Millisecond
-	allowable := timeout + 10*time.Millisecond
-	abortAfter := timeout + 100*time.Millisecond
+			timeout := time.Millisecond
+			allowable := timeout + 10*time.Millisecond
 
-	start := time.Now()
+			_, _, err := exchange(m, addrstr, timeout)
+			if err == nil {
+				t.Errorf("no timeout using Client.%s", name)
+			}
 
-	go func() {
+			length := time.Since(start)
+			if length > allowable {
+				t.Errorf("exchange took longer %v than specified Timeout %v", length, allowable)
+			}
+		})
+	}
+	runTest("Exchange", func(m *Msg, addr string, timeout time.Duration) (*Msg, time.Duration, error) {
 		c := &Client{Timeout: timeout}
-		_, _, err := c.Exchange(m, addrstr)
-		if err == nil {
-			t.Error("no timeout using Client.Exchange")
-		}
-		done <- struct{}{}
-	}()
-
-	go func() {
+		return c.Exchange(m, addr)
+	})
+	runTest("ExchangeContext", func(m *Msg, addr string, timeout time.Duration) (*Msg, time.Duration, error) {
 		ctx, cancel := context.WithTimeout(context.Background(), timeout)
 		defer cancel()
-		c := &Client{}
-		_, _, err := c.ExchangeContext(ctx, m, addrstr)
-		if err == nil {
-			t.Error("no timeout using Client.ExchangeContext")
-		}
-		done <- struct{}{}
-	}()
 
-	// Wait for both the Exchange and ExchangeContext tests to be done.
-	for i := 0; i < 2; i++ {
-		select {
-		case <-done:
-		case <-time.After(abortAfter):
-		}
-	}
-
-	length := time.Since(start)
-
-	if length > allowable {
-		t.Errorf("exchange took longer %v than specified Timeout %v", length, allowable)
-	}
+		return new(Client).ExchangeContext(ctx, m, addrstr)
+	})
 }
 
 // Check that responses from deduplicated requests aren't shared between callers