plugin/loadbalance: add parse and tests (#1947)

Automatically submitted.

4 files changed, 71 insertions, 16 deletions

diff --git a/plugin/loadbalance/README.md b/plugin/loadbalance/README.md
index 4c075bb9d..fed9c906f 100644
--- a/plugin/loadbalance/README.md
+++ b/plugin/loadbalance/README.md
@@ -2,12 +2,12 @@
 
 ## Name
 
-*loadbalance* - randomize the order of A and AAAA records.
+*loadbalance* - randomize the order of A, AAAA and MX records.
 
 ## Description
 
-The *loadbalance* will act as a round-robin DNS loadbalancer by randomizing the order of A and AAAA
-records in the answer.
+The *loadbalance* will act as a round-robin DNS loadbalancer by randomizing the order of A, AAAA,
+and MX records in the answer.
  
 See [Wikipedia](https://en.wikipedia.org/wiki/Round-robin_DNS) about the pros and cons on this
 setup. It will take care to sort any CNAMEs before any address records, because some stub resolver
diff --git a/plugin/loadbalance/loadbalance.go b/plugin/loadbalance/loadbalance.go
index 031f841f4..39e70368d 100644
--- a/plugin/loadbalance/loadbalance.go
+++ b/plugin/loadbalance/loadbalance.go
@@ -1,14 +1,12 @@
-// Package loadbalance shuffles A and AAAA records.
+// Package loadbalance shuffles A, AAAA and MX records.
 package loadbalance
 
 import (
 	"github.com/miekg/dns"
 )
 
-// RoundRobinResponseWriter is a response writer that shuffles A and AAAA records.
-type RoundRobinResponseWriter struct {
-	dns.ResponseWriter
-}
+// RoundRobinResponseWriter is a response writer that shuffles A, AAAA and MX records.
+type RoundRobinResponseWriter struct{ dns.ResponseWriter }
 
 // WriteMsg implements the dns.ResponseWriter interface.
 func (r *RoundRobinResponseWriter) WriteMsg(res *dns.Msg) error {
@@ -77,9 +75,3 @@ func (r *RoundRobinResponseWriter) Write(buf []byte) (int, error) {
 	n, err := r.ResponseWriter.Write(buf)
 	return n, err
 }
-
-// Hijack implements the dns.ResponseWriter interface.
-func (r *RoundRobinResponseWriter) Hijack() {
-	r.ResponseWriter.Hijack()
-	return
-}
diff --git a/plugin/loadbalance/setup.go b/plugin/loadbalance/setup.go
index 38dce6308..9d52a3cbe 100644
--- a/plugin/loadbalance/setup.go
+++ b/plugin/loadbalance/setup.go
@@ -1,6 +1,8 @@
 package loadbalance
 
 import (
+	"fmt"
+
 	"github.com/coredns/coredns/core/dnsserver"
 	"github.com/coredns/coredns/plugin"
 	clog "github.com/coredns/coredns/plugin/pkg/log"
@@ -18,8 +20,9 @@ func init() {
 }
 
 func setup(c *caddy.Controller) error {
-	for c.Next() {
-		// TODO(miek): block and option parsing
+	err := parse(c)
+	if err != nil {
+		return plugin.Error("loadbalance", err)
 	}
 
 	dnsserver.GetConfig(c).AddPlugin(func(next plugin.Handler) plugin.Handler {
@@ -28,3 +31,20 @@ func setup(c *caddy.Controller) error {
 
 	return nil
 }
+
+func parse(c *caddy.Controller) error {
+	for c.Next() {
+		args := c.RemainingArgs()
+		switch len(args) {
+		case 0:
+			return nil
+		case 1:
+			if args[0] != "round_robin" {
+				return fmt.Errorf("unknown policy: %s", args[0])
+
+			}
+			return nil
+		}
+	}
+	return c.ArgErr()
+}
diff --git a/plugin/loadbalance/setup_test.go b/plugin/loadbalance/setup_test.go
new file mode 100644
index 000000000..67673a72c
--- /dev/null
+++ b/plugin/loadbalance/setup_test.go
@@ -0,0 +1,43 @@
+package loadbalance
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/mholt/caddy"
+)
+
+func TestSetup(t *testing.T) {
+	tests := []struct {
+		input              string
+		shouldErr          bool
+		expectedPolicy     string
+		expectedErrContent string // substring from the expected error. Empty for positive cases.
+	}{
+		// positive
+		{`loadbalance`, false, "round_robin", ""},
+		{`loadbalance round_robin`, false, "round_robin", ""},
+		// negative
+		{`loadbalance fleeb`, true, "", "unknown policy"},
+		{`loadbalance a b`, true, "", "argument count or unexpected line"},
+	}
+
+	for i, test := range tests {
+		c := caddy.NewTestController("dns", test.input)
+		err := parse(c)
+
+		if test.shouldErr && err == nil {
+			t.Errorf("Test %d: Expected error but found %s for input %s", i, err, test.input)
+		}
+
+		if err != nil {
+			if !test.shouldErr {
+				t.Errorf("Test %d: Expected no error but found one for input %s. Error was: %v", i, test.input, err)
+			}
+
+			if !strings.Contains(err.Error(), test.expectedErrContent) {
+				t.Errorf("Test %d: Expected error to contain: %v, found error: %v, input: %s", i, test.expectedErrContent, err, test.input)
+			}
+		}
+	}
+}