From 77a3bfd606cf78a16b5a8f262c1db2acbd6b82b3 Mon Sep 17 00:00:00 2001
From: Armon Dadgar <armon.dadgar@gmail.com>
Date: Thu, 8 Jan 2015 11:58:32 -0800
Subject: [PATCH] agent: Adding test for DNS lookup with service address

---
 command/agent/dns_test.go | 69 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 69 insertions(+)

diff --git a/command/agent/dns_test.go b/command/agent/dns_test.go
index 4c713bb4c2..dbd7fad745 100644
--- a/command/agent/dns_test.go
+++ b/command/agent/dns_test.go
@@ -430,6 +430,7 @@ func TestDNS_ReverseLookup_IPV6(t *testing.T) {
 		t.Fatalf("Bad: %#v", ptrRec)
 	}
 }
+
 func TestDNS_ServiceLookup(t *testing.T) {
 	dir, srv := makeDNSServer(t)
 	defer os.RemoveAll(dir)
@@ -497,6 +498,74 @@ func TestDNS_ServiceLookup(t *testing.T) {
 	}
 }
 
+func TestDNS_ServiceLookup_ServiceAddress(t *testing.T) {
+	dir, srv := makeDNSServer(t)
+	defer os.RemoveAll(dir)
+	defer srv.agent.Shutdown()
+
+	testutil.WaitForLeader(t, srv.agent.RPC, "dc1")
+
+	// Register node
+	args := &structs.RegisterRequest{
+		Datacenter: "dc1",
+		Node:       "foo",
+		Address:    "127.0.0.1",
+		Service: &structs.NodeService{
+			Service: "db",
+			Tags:    []string{"master"},
+			Address: "127.0.0.2",
+			Port:    12345,
+		},
+	}
+
+	var out struct{}
+	if err := srv.agent.RPC("Catalog.Register", args, &out); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	m := new(dns.Msg)
+	m.SetQuestion("db.service.consul.", dns.TypeSRV)
+
+	c := new(dns.Client)
+	addr, _ := srv.agent.config.ClientListener("", srv.agent.config.Ports.DNS)
+	in, _, err := c.Exchange(m, addr.String())
+	if err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	if len(in.Answer) != 1 {
+		t.Fatalf("Bad: %#v", in)
+	}
+
+	srvRec, ok := in.Answer[0].(*dns.SRV)
+	if !ok {
+		t.Fatalf("Bad: %#v", in.Answer[0])
+	}
+	if srvRec.Port != 12345 {
+		t.Fatalf("Bad: %#v", srvRec)
+	}
+	if srvRec.Target != "foo.node.dc1.consul." {
+		t.Fatalf("Bad: %#v", srvRec)
+	}
+	if srvRec.Hdr.Ttl != 0 {
+		t.Fatalf("Bad: %#v", in.Answer[0])
+	}
+
+	aRec, ok := in.Extra[0].(*dns.A)
+	if !ok {
+		t.Fatalf("Bad: %#v", in.Extra[0])
+	}
+	if aRec.Hdr.Name != "foo.node.dc1.consul." {
+		t.Fatalf("Bad: %#v", in.Extra[0])
+	}
+	if aRec.A.String() != "127.0.0.2" {
+		t.Fatalf("Bad: %#v", in.Extra[0])
+	}
+	if aRec.Hdr.Ttl != 0 {
+		t.Fatalf("Bad: %#v", in.Extra[0])
+	}
+}
+
 func TestDNS_CaseInsensitiveServiceLookup(t *testing.T) {
 	dir, srv := makeDNSServer(t)
 	defer os.RemoveAll(dir)