From 5089561b1bdb30f8c0022fd43a0d9c8a87eeea3d Mon Sep 17 00:00:00 2001
From: Ryan Uber <ru@ryanuber.com>
Date: Fri, 1 Jul 2016 09:46:26 -0700
Subject: [PATCH] agent: test that DNS passes the agent data through

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

diff --git a/command/agent/dns_test.go b/command/agent/dns_test.go
index 863a0bfef..1182d9172 100644
--- a/command/agent/dns_test.go
+++ b/command/agent/dns_test.go
@@ -3166,3 +3166,37 @@ func TestDNS_InvalidQueries(t *testing.T) {
 		}
 	}
 }
+
+func TestDNS_PreparedQuery_AgentSource(t *testing.T) {
+	dir, srv := makeDNSServer(t)
+	defer os.RemoveAll(dir)
+	defer srv.agent.Shutdown()
+
+	testutil.WaitForLeader(t, srv.agent.RPC, "dc1")
+
+	m := MockPreparedQuery{}
+	if err := srv.agent.InjectEndpoint("PreparedQuery", &m); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	m.executeFn = func(args *structs.PreparedQueryExecuteRequest, reply *structs.PreparedQueryExecuteResponse) error {
+		// Check that the agent inserted its self-name and datacenter to
+		// the RPC request body.
+		if args.Agent.Datacenter != srv.agent.config.Datacenter ||
+			args.Agent.Node != srv.agent.config.NodeName {
+			t.Fatalf("bad: %#v", args.Agent)
+		}
+		return nil
+	}
+
+	{
+		m := new(dns.Msg)
+		m.SetQuestion("foo.query.consul.", dns.TypeSRV)
+
+		c := new(dns.Client)
+		addr, _ := srv.agent.config.ClientListener("", srv.agent.config.Ports.DNS)
+		if _, _, err := c.Exchange(m, addr.String()); err != nil {
+			t.Fatalf("err: %v", err)
+		}
+	}
+}