Added ratelimit whitelist and tests

coredns_plugin/ratelimit/ratelimit_test.go

@@ -0,0 +1,82 @@
+package ratelimit
+
+import (
+	"testing"
+
+	"github.com/mholt/caddy"
+)
+
+func TestSetup(t *testing.T) {
+	for i, testcase := range []struct {
+		config  string
+		failing bool
+	}{
+		{`ratelimit`, false},
+		{`ratelimit 100`, false},
+		{`ratelimit { 
+					whitelist 127.0.0.1
+				}`, false},
+		{`ratelimit 50 {
+					whitelist 127.0.0.1 176.103.130.130
+				}`, false},
+		{`ratelimit test`, true},
+	} {
+		c := caddy.NewTestController("dns", testcase.config)
+		err := setup(c)
+		if err != nil {
+			if !testcase.failing {
+				t.Fatalf("Test #%d expected no errors, but got: %v", i, err)
+			}
+			continue
+		}
+		if testcase.failing {
+			t.Fatalf("Test #%d expected to fail but it didn't", i)
+		}
+	}
+}
+
+func TestRatelimiting(t *testing.T) {
+
+	// rate limit is 1 per sec
+	c := caddy.NewTestController("dns", `ratelimit 1`)
+	p, err := setupPlugin(c)
+
+	if err != nil {
+		t.Fatal("Failed to initialize the plugin")
+	}
+
+	allowed, err := p.allowRequest("127.0.0.1")
+
+	if err != nil || !allowed {
+		t.Fatal("First request must have been allowed")
+	}
+
+	allowed, err = p.allowRequest("127.0.0.1")
+
+	if err != nil || allowed {
+		t.Fatal("Second request must have been ratelimited")
+	}
+}
+
+func TestWhitelist(t *testing.T) {
+
+	// rate limit is 1 per sec
+	c := caddy.NewTestController("dns", `ratelimit 1 { whitelist 127.0.0.1 }`)
+	p, err := setupPlugin(c)
+
+	if err != nil {
+		t.Fatal("Failed to initialize the plugin")
+	}
+
+	allowed, err := p.allowRequest("127.0.0.1")
+
+	if err != nil || !allowed {
+		t.Fatal("First request must have been allowed")
+	}
+
+	allowed, err = p.allowRequest("127.0.0.1")
+
+	if err != nil || !allowed {
+		t.Fatal("Second request must have been allowed due to whitelist")
+	}
+}