From 27a820d611fe76ffb997bf259a076c063efebd1c Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Tue, 5 May 2015 14:14:41 -0700
Subject: [PATCH] agent: Adding test for DNS enable_truncate

---
 command/agent/dns.go      |  3 ++-
 command/agent/dns_test.go | 45 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/command/agent/dns.go b/command/agent/dns.go
index cc49291787..f4d8bb8891 100644
--- a/command/agent/dns.go
+++ b/command/agent/dns.go
@@ -492,8 +492,9 @@ RPC:
 	// If the network is not TCP, restrict the number of responses
 	if network != "tcp" && len(out.Nodes) > maxServiceResponses {
 		out.Nodes = out.Nodes[:maxServiceResponses]
+
 		// Flag that there are more records to return in the UDP response
-		if d.config.EnableTruncate == true {
+		if d.config.EnableTruncate {
 			resp.Truncated = true
 		}
 	}
diff --git a/command/agent/dns_test.go b/command/agent/dns_test.go
index 90f7f58264..ec584a1cd2 100644
--- a/command/agent/dns_test.go
+++ b/command/agent/dns_test.go
@@ -1217,6 +1217,51 @@ func TestDNS_ServiceLookup_Randomize(t *testing.T) {
 	}
 }
 
+func TestDNS_ServiceLookup_Truncate(t *testing.T) {
+	config := &DNSConfig{
+		EnableTruncate: true,
+	}
+	dir, srv := makeDNSServer(t, config, nil)
+	defer os.RemoveAll(dir)
+	defer srv.agent.Shutdown()
+
+	testutil.WaitForLeader(t, srv.agent.RPC, "dc1")
+
+	// Register nodes
+	for i := 0; i < 3*maxServiceResponses; i++ {
+		args := &structs.RegisterRequest{
+			Datacenter: "dc1",
+			Node:       fmt.Sprintf("foo%d", i),
+			Address:    fmt.Sprintf("127.0.0.%d", i+1),
+			Service: &structs.NodeService{
+				Service: "web",
+				Port:    8000,
+			},
+		}
+
+		var out struct{}
+		if err := srv.agent.RPC("Catalog.Register", args, &out); err != nil {
+			t.Fatalf("err: %v", err)
+		}
+	}
+
+	// Ensure the response is randomized each time.
+	m := new(dns.Msg)
+	m.SetQuestion("web.service.consul.", dns.TypeANY)
+
+	addr, _ := srv.agent.config.ClientListener("", srv.agent.config.Ports.DNS)
+	c := new(dns.Client)
+	in, _, err := c.Exchange(m, addr.String())
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	// Check for the truncate bit
+	if !in.Truncated {
+		t.Fatalf("should have truncate bit")
+	}
+}
+
 func TestDNS_ServiceLookup_CNAME(t *testing.T) {
 	recursor := makeRecursor(t, []dns.RR{
 		dnsCNAME("www.google.com", "google.com"),