Add more thorough testing for only_passing behavior

command/agent/dns_test.go

@@ -954,6 +954,26 @@ func TestDNS_ServiceLookup_FilterCritical(t *testing.T) {
 		t.Fatalf("err: %v", err)
 	}
 
+	args5 := &structs.RegisterRequest{
+		Datacenter: "dc1",
+		Node:       "quux",
+		Address:    "127.0.0.4",
+		Service: &structs.NodeService{
+			Service: "db",
+			Tags:    []string{"master"},
+			Port:    12345,
+		},
+		Check: &structs.HealthCheck{
+			CheckID:   "db",
+			Name:      "db",
+			ServiceID: "db",
+			Status:    structs.HealthWarning,
+		},
+	}
+	if err := srv.agent.RPC("Catalog.Register", args5, &out); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
 	m := new(dns.Msg)
 	m.SetQuestion("db.service.consul.", dns.TypeANY)
 
@@ -964,14 +984,137 @@ func TestDNS_ServiceLookup_FilterCritical(t *testing.T) {
 		t.Fatalf("err: %v", err)
 	}
 
-	// Should get no answer since we are failing!
+	// Only 4 and 5 are not failing, so we should get 2 answers
+	if len(in.Answer) != 2 {
+		t.Fatalf("Bad: %#v", in)
+	}
+
+	ips := make(map[string]bool)
+	for _, resp := range in.Answer {
+		aRec := resp.(*dns.A)
+		ips[aRec.A.String()] = true
+	}
+
+	if !ips["127.0.0.3"] {
+		t.Fatalf("Bad: %#v should contain 127.0.0.3 (state healthy)", in)
+	}
+	if !ips["127.0.0.4"] {
+		t.Fatalf("Bad: %#v should contain 127.0.0.4 (state warning)", in)
+	}
+}
+
+func TestDNS_ServiceLookup_OnlyPassing(t *testing.T) {
+	dir, srv := makeDNSServerConfig(t, &DNSConfig{OnlyPassing: true})
+	defer os.RemoveAll(dir)
+	defer srv.agent.Shutdown()
+
+	testutil.WaitForLeader(t, srv.agent.RPC, "dc1")
+
+	// Register nodes
+	args := &structs.RegisterRequest{
+		Datacenter: "dc1",
+		Node:       "foo",
+		Address:    "127.0.0.1",
+		Service: &structs.NodeService{
+			Service: "db",
+			Tags:    []string{"master"},
+			Port:    12345,
+		},
+		Check: &structs.HealthCheck{
+			CheckID:   "db",
+			Name:      "db",
+			ServiceID: "db",
+			Status:    structs.HealthPassing,
+		},
+	}
+
+	var out struct{}
+	if err := srv.agent.RPC("Catalog.Register", args, &out); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	args2 := &structs.RegisterRequest{
+		Datacenter: "dc1",
+		Node:       "bar",
+		Address:    "127.0.0.2",
+		Service: &structs.NodeService{
+			Service: "db",
+			Tags:    []string{"master"},
+			Port:    12345,
+		},
+		Check: &structs.HealthCheck{
+			CheckID:   "db",
+			Name:      "db",
+			ServiceID: "db",
+			Status:    structs.HealthWarning,
+		},
+	}
+
+	if err := srv.agent.RPC("Catalog.Register", args2, &out); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	args3 := &structs.RegisterRequest{
+		Datacenter: "dc1",
+		Node:       "baz",
+		Address:    "127.0.0.3",
+		Service: &structs.NodeService{
+			Service: "db",
+			Tags:    []string{"master"},
+			Port:    12345,
+		},
+		Check: &structs.HealthCheck{
+			CheckID:   "db",
+			Name:      "db",
+			ServiceID: "db",
+			Status:    structs.HealthCritical,
+		},
+	}
+
+	if err := srv.agent.RPC("Catalog.Register", args3, &out); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	args4 := &structs.RegisterRequest{
+		Datacenter: "dc1",
+		Node:       "quux",
+		Address:    "127.0.0.4",
+		Service: &structs.NodeService{
+			Service: "db",
+			Tags:    []string{"master"},
+			Port:    12345,
+		},
+		Check: &structs.HealthCheck{
+			CheckID:   "db",
+			Name:      "db",
+			ServiceID: "db",
+			Status:    structs.HealthUnknown,
+		},
+	}
+
+	if err := srv.agent.RPC("Catalog.Register", args4, &out); err != nil {
+		t.Fatalf("err: %v", err)
+	}
+
+	m := new(dns.Msg)
+	m.SetQuestion("db.service.consul.", dns.TypeANY)
+
+	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)
+	}
+
+	// Only 1 is passing, so we should only get 1 answer
 	if len(in.Answer) != 1 {
 		t.Fatalf("Bad: %#v", in)
 	}
 
 	resp := in.Answer[0]
 	aRec := resp.(*dns.A)
-	if aRec.A.String() != "127.0.0.3" {
+
+	if aRec.A.String() != "127.0.0.1" {
 		t.Fatalf("Bad: %#v", in.Answer[0])
 	}
 }