From e7a401841e212a962c3cd423345f907b50bd3bc8 Mon Sep 17 00:00:00 2001
From: Alex Sergeyev <asergeyev@dyn.com>
Date: Fri, 29 Aug 2014 08:59:19 -0400
Subject: [PATCH] Fixed client tests and removed TODO about timeouts

---
 README.md      |  1 -
 client_test.go | 52 ++++++++++++++++++++++++++++++++++----------------
 2 files changed, 36 insertions(+), 17 deletions(-)

diff --git a/README.md b/README.md
index 11d46849..b259d3c5 100644
--- a/README.md
+++ b/README.md
@@ -133,4 +133,3 @@ Example programs can be found in the `github.com/miekg/exdns` repository.
 * CAA parsing is broken;
 * Replies with TC bit are not parsed to the end.
 * SIG(0)
-* Make tests run faster by removing the time.Sleep()s
diff --git a/client_test.go b/client_test.go
index bb1c714c..b6b206c2 100644
--- a/client_test.go
+++ b/client_test.go
@@ -9,27 +9,25 @@ import (
 	"time"
 )
 
-func newTestServer(t *testing.T) {
-	// Defined in server_test.go
+func TestClientSync(t *testing.T) {
 	HandleFunc("miek.nl.", HelloServer)
 	HandleFunc("example.com.", AnotherHelloServer)
-	go func() {
-		err := ListenAndServe(":8063", "udp", nil)
-		if err != nil {
-			t.Log("ListenAndServe: ", err.Error())
-			t.Fatal()
-		}
-	}()
-	time.Sleep(4e8)
-}
 
-func TestClientSync(t *testing.T) {
+	s, addrstr, err := RunLocalUDPServer("127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("Unable to run test server: %s", err)
+	}
+	defer s.Shutdown()
+
 	m := new(Msg)
 	m.SetQuestion("miek.nl.", TypeSOA)
 
 	c := new(Client)
-	r, _, _ := c.Exchange(m, "127.0.0.1:6053")
-
+	r, _, e := c.Exchange(m, addrstr)
+	if e != nil {
+		t.Logf("failed to exchange: %s", e.Error())
+		t.Fail()
+	}
 	if r != nil && r.Rcode != RcodeSuccess {
 		t.Log("failed to get an valid answer")
 		t.Fail()
@@ -38,13 +36,26 @@ func TestClientSync(t *testing.T) {
 }
 
 func TestClientEDNS0(t *testing.T) {
+	HandleFunc("miek.nl.", HelloServer)
+	HandleFunc("example.com.", AnotherHelloServer)
+
+	s, addrstr, err := RunLocalUDPServer("127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("Unable to run test server: %s", err)
+	}
+	defer s.Shutdown()
+
 	m := new(Msg)
 	m.SetQuestion("miek.nl.", TypeDNSKEY)
 
 	m.SetEdns0(2048, true)
 
 	c := new(Client)
-	r, _, _ := c.Exchange(m, "127.0.0.1:6053")
+	r, _, e := c.Exchange(m, addrstr)
+	if e != nil {
+		t.Logf("failed to exchange: %s", e.Error())
+		t.Fail()
+	}
 
 	if r != nil && r.Rcode != RcodeSuccess {
 		t.Log("failed to get an valid answer")
@@ -54,6 +65,15 @@ func TestClientEDNS0(t *testing.T) {
 }
 
 func TestSingleSingleInflight(t *testing.T) {
+	HandleFunc("miek.nl.", HelloServer)
+	HandleFunc("example.com.", AnotherHelloServer)
+
+	s, addrstr, err := RunLocalUDPServer("127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("Unable to run test server: %s", err)
+	}
+	defer s.Shutdown()
+
 	m := new(Msg)
 	m.SetQuestion("miek.nl.", TypeDNSKEY)
 
@@ -63,7 +83,7 @@ func TestSingleSingleInflight(t *testing.T) {
 	ch := make(chan time.Duration)
 	for i := 0; i < nr; i++ {
 		go func() {
-			_, rtt, _ := c.Exchange(m, "127.0.0.1:6053")
+			_, rtt, _ := c.Exchange(m, addrstr)
 			ch <- rtt
 		}()
 	}