From 11d93b1883cc5c0cd5ce36c4733ea2873b10de9e Mon Sep 17 00:00:00 2001
From: Tom Wilkie <tom@weave.works>
Date: Wed, 30 Sep 2015 15:51:07 +0000
Subject: [PATCH] Fix races in TestRecursiveCompress

---
 nameserver/dns_test.go | 18 ++++++++++--------
 1 file changed, 10 insertions(+), 8 deletions(-)

diff --git a/nameserver/dns_test.go b/nameserver/dns_test.go
index 6bc9da7a..13bc6153 100644
--- a/nameserver/dns_test.go
+++ b/nameserver/dns_test.go
@@ -15,7 +15,7 @@ import (
 	"github.com/weaveworks/weave/router"
 )
 
-func startServer(t *testing.T) (*DNSServer, *Nameserver, int, int) {
+func startServer(t *testing.T, upstream *dns.ClientConfig) (*DNSServer, *Nameserver, int, int) {
 	peername, err := router.PeerNameFromString("00:00:00:02:00:00")
 	require.Nil(t, err)
 	nameserver := New(peername, nil, nil, "")
@@ -23,13 +23,16 @@ func startServer(t *testing.T) (*DNSServer, *Nameserver, int, int) {
 	require.Nil(t, err)
 	udpPort := dnsserver.servers[0].PacketConn.LocalAddr().(*net.UDPAddr).Port
 	tcpPort := dnsserver.servers[1].Listener.Addr().(*net.TCPAddr).Port
+	if upstream != nil {
+		dnsserver.upstream = upstream
+	}
 	go dnsserver.ActivateAndServe()
 	return dnsserver, nameserver, udpPort, tcpPort
 }
 
 func TestTruncation(t *testing.T) {
 	//common.SetLogLevel("debug")
-	dnsserver, nameserver, udpPort, tcpPort := startServer(t)
+	dnsserver, nameserver, udpPort, tcpPort := startServer(t, nil)
 	defer dnsserver.Stop()
 
 	// Add 100 mappings to nameserver
@@ -122,9 +125,9 @@ func TestRecursiveCompress(t *testing.T) {
 	require.True(t, response.Len() <= maxSize)
 
 	// A dns server that returns the above response
-	var gotRequest = false
+	var gotRequest = make(chan struct{}, 1)
 	handleRecursive := func(w dns.ResponseWriter, req *dns.Msg) {
-		gotRequest = true
+		gotRequest <- struct{}{}
 		require.Equal(t, req.Question[0].Name, hostname)
 		response.SetReply(req)
 		err := w.WriteMsg(&response)
@@ -140,14 +143,13 @@ func TestRecursiveCompress(t *testing.T) {
 	defer udpServer.Shutdown()
 
 	// The weavedns server, pointed at the above server
-	dnsserver, _, udpPort, _ := startServer(t)
-	dnsserver.upstream = &dns.ClientConfig{
+	dnsserver, _, udpPort, _ := startServer(t, &dns.ClientConfig{
 		Servers:  []string{"127.0.0.1"},
 		Port:     strconv.Itoa(udpServerPort),
 		Ndots:    1,
 		Timeout:  5,
 		Attempts: 2,
-	}
+	})
 	defer dnsserver.Stop()
 
 	// Now do lookup, check its what we expected.
@@ -165,6 +167,6 @@ func TestRecursiveCompress(t *testing.T) {
 	c := new(dns.Client)
 	res, _, err := c.Exchange(req, fmt.Sprintf("127.0.0.1:%d", udpPort))
 	require.Nil(t, err)
-	require.True(t, gotRequest)
+	require.True(t, len(gotRequest) > 0)
 	require.True(t, res.Len() > maxSize)
 }