From 31a8ac3219907138303701604b7bf3bcbff4b249 Mon Sep 17 00:00:00 2001
From: johncming <johncming@yahoo.com>
Date: Mon, 23 Sep 2019 19:26:12 +0800
Subject: [PATCH] discovery/dns: add test case for SDConfig.UnmarshalYAML.
 (#6035)

* discovery/dns: Add code coverage.

Signed-off-by: johncming <johncming@yahoo.com>

* discovery/dns: add test case for SDConfig.UnmarshalYAML.

Signed-off-by: johncming <johncming@yahoo.com>
---
 discovery/dns/dns_test.go | 94 ++++++++++++++++++++++++++++++++++++++-
 1 file changed, 93 insertions(+), 1 deletion(-)

diff --git a/discovery/dns/dns_test.go b/discovery/dns/dns_test.go
index 8225ddf81..c5d6f361d 100644
--- a/discovery/dns/dns_test.go
+++ b/discovery/dns/dns_test.go
@@ -22,8 +22,9 @@ import (
 
 	"github.com/go-kit/kit/log"
 	"github.com/miekg/dns"
-	"github.com/prometheus/common/model"
+	"gopkg.in/yaml.v2"
 
+	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/discovery/targetgroup"
 	"github.com/prometheus/prometheus/util/testutil"
 )
@@ -103,6 +104,7 @@ func TestDNS(t *testing.T) {
 			name: "SRV record query",
 			config: SDConfig{
 				Names:           []string{"_mysql._tcp.db.example.com."},
+				Type:            "SRV",
 				RefreshInterval: model.Duration(time.Minute),
 			},
 			lookup: func(name string, qtype uint16, logger log.Logger) (*dns.Msg, error) {
@@ -178,3 +180,93 @@ func TestDNS(t *testing.T) {
 		})
 	}
 }
+
+func TestSDConfigUnmarshalYAML(t *testing.T) {
+	marshal := func(c SDConfig) []byte {
+		d, err := yaml.Marshal(c)
+		if err != nil {
+			panic(err)
+		}
+		return d
+	}
+
+	unmarshal := func(d []byte) func(interface{}) error {
+		return func(o interface{}) error {
+			return yaml.Unmarshal(d, o)
+		}
+	}
+
+	cases := []struct {
+		name      string
+		input     SDConfig
+		expectErr bool
+	}{
+		{
+			name: "valid srv",
+			input: SDConfig{
+				Names: []string{"a.example.com", "b.example.com"},
+				Type:  "SRV",
+			},
+			expectErr: false,
+		},
+		{
+			name: "valid a",
+			input: SDConfig{
+				Names: []string{"a.example.com", "b.example.com"},
+				Type:  "A",
+				Port:  5300,
+			},
+			expectErr: false,
+		},
+		{
+			name: "valid aaaa",
+			input: SDConfig{
+				Names: []string{"a.example.com", "b.example.com"},
+				Type:  "AAAA",
+				Port:  5300,
+			},
+			expectErr: false,
+		},
+		{
+			name: "invalid a without port",
+			input: SDConfig{
+				Names: []string{"a.example.com", "b.example.com"},
+				Type:  "A",
+			},
+			expectErr: true,
+		},
+		{
+			name: "invalid aaaa without port",
+			input: SDConfig{
+				Names: []string{"a.example.com", "b.example.com"},
+				Type:  "AAAA",
+			},
+			expectErr: true,
+		},
+		{
+			name: "invalid empty names",
+			input: SDConfig{
+				Names: []string{},
+				Type:  "AAAA",
+			},
+			expectErr: true,
+		},
+		{
+			name: "invalid unkown dns type",
+			input: SDConfig{
+				Names: []string{"a.example.com", "b.example.com"},
+				Type:  "PTR",
+			},
+			expectErr: true,
+		},
+	}
+
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
+			var config SDConfig
+			d := marshal(c.input)
+			err := config.UnmarshalYAML(unmarshal(d))
+			testutil.Equals(t, c.expectErr, err != nil)
+		})
+	}
+}