Make HTTP middleware easier to use

server/router/router_test.go

Index: server/router/router_test.go
diff --git a/server/router/router_test.go b/server/router/router_test.go
new file mode 100644
index 0000000000000000000000000000000000000000..f21bf50be19fd7f0302d28dbf742732ec0342ef3
--- /dev/null
+++ b/server/router/router_test.go
@@ -0,0 +1,323 @@
+// Copyright 2016 The LUCI Authors. All rights reserved.
+// Use of this source code is governed under the Apache License, Version 2.0
+// that can be found in the LICENSE file.
+
+package router
+
+import (
+	"io"
+	"io/ioutil"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"golang.org/x/net/context"
+
+	assert "github.com/smartystreets/assertions"
+	. "github.com/smartystreets/goconvey/convey"
+)
+
+func TestNew(t *testing.T) {
+	t.Parallel()
+	Convey("Initializes non-nil httprouter.Router", t, func() {
+		So(New().hrouter, ShouldNotBeNil)
+	})
+}
+
+func TestUse(t *testing.T) {
+	t.Parallel()
+	a := func(_ *Context, _ Handler) {}
+	b := func(_ *Context, _ Handler) {}
+
+	Convey("Use appends middleware", t, func() {
+		r := New()
+		So(len(r.middleware), ShouldEqual, 0)
+		r.Use(MiddlewareChain{a, b})
+		So(len(r.middleware), ShouldEqual, 2)
+	})
+}
+
+func TestGroup(t *testing.T) {
+	t.Parallel()
+	a := func(_ *Context, _ Handler) {}
+	b := func(_ *Context, _ Handler) {}
+
+	Convey("Group derives from original router", t, func() {
+		r1 := New()
+		r1.Use(MiddlewareChain{a, b})
+		r2 := r1.Group("/foo")
+		So(r1.hrouter, ShouldPointTo, r2.hrouter)
+		So(&r1.middleware[0], ShouldNotPointTo, &r2.middleware[0])
+		So(len(r1.middleware), ShouldEqual, len(r2.middleware))
+	})
+}
+
+func TestHandle(t *testing.T) {
+	t.Parallel()
+	a := func(_ *Context, _ Handler) {}
+	b := func(_ *Context, _ Handler) {}
+	c := func(_ *Context, _ Handler) {}
+	handler := func(_ *Context) {}
+
+	Convey("No side-effects on r.middleware", t, func() {
+		r := New()
+		So(len(r.middleware), ShouldEqual, 0)
+		r.Handle("GET", "/bar", MiddlewareChain{b, c}, handler)
+		So(len(r.middleware), ShouldEqual, 0)
+	})
+
+	Convey("No side-effects on r.middleware", t, func() {
+		r := New()
+		r.Use(MiddlewareChain{a})
+		So(len(r.middleware), ShouldEqual, 1)
+		r.Handle("GET", "/bar", MiddlewareChain{b, c}, handler)
+		So(len(r.middleware), ShouldEqual, 1)
+	})
+}
+
+func TestRunMiddlewareChain(t *testing.T) {
+	t.Parallel()
+	a := func(c *Context, next Handler) {
+		io.WriteString(c, "1")
+		next(c)
+		io.WriteString(c, "7")
+	}
+	b := func(c *Context, next Handler) {
+		io.WriteString(c, "2")
+		next(c)
+		io.WriteString(c, "6")
+	}
+	c := func(c *Context, next Handler) {
+		io.WriteString(c, "3")
+		next(c)
+	}
+	d := func(c *Context, next Handler) {
+		next(c)
+		io.WriteString(c, "5")
+	}
+	stop := func(_ *Context, _ Handler) {}
+	handler := func(c *Context) {
+		io.WriteString(c, "4")
+	}
+
+	Convey("Handles length 0 chains", t, func() {
+		rec := httptest.NewRecorder()
+		var m MiddlewareChain
+		n := MiddlewareChain{}
+		run(&Context{writer: rec}, m, n, handler)
+		So(rec.Body.String(), ShouldResemble, "4")
+	})
+
+	Convey("Executes in order", t, func() {
+		rec := httptest.NewRecorder()
+		m := MiddlewareChain{a, b, c}
+		n := MiddlewareChain{d}
+		run(&Context{writer: rec}, m, n, handler)
+		So(rec.Body.String(), ShouldResemble, "1234567")
+	})
+
+	Convey("Not calling next does not execute upcoming handlers", t, func() {
+		rec := httptest.NewRecorder()
+		mc := MiddlewareChain{a, stop, b}
+		mc.Run(&Context{writer: rec}, handler)
+		So(rec.Body.String(), ShouldResemble, "17")
+	})
+}
+
+func TestMergeMiddlewareChains(t *testing.T) {
+	t.Parallel()
+	a := func(_ *Context, _ Handler) {}
+	b := func(_ *Context, _ Handler) {}
+	c := func(_ *Context, _ Handler) {}
+	d := func(_ *Context, _ Handler) {}
+
+	Convey("Merged length is sum of lengths", t, func() {
+		mc1 := MiddlewareChain{a, b}
+		mc2 := MiddlewareChain{c}
+		merged := MergeMiddlewareChains(mc1, mc2)
+		So(len(merged), ShouldEqual, 3)
+	})
+
+	Convey("Merged slice points to separate array", t, func() {
+		mc1 := MiddlewareChain{a, b}
+		mc2 := MiddlewareChain{c, d}
+		merged := MergeMiddlewareChains(mc1, mc2)
+		So(&merged[0], ShouldNotPointTo, &mc1[0])
+	})
+
+	Convey("Functions in merged slice are same", t, func() {
+		mc1 := MiddlewareChain{a, b}
+		mc2 := MiddlewareChain{c, d}
+		merged := MergeMiddlewareChains(mc1, mc2)
+		So(merged[0], ShouldEqual, mc1[0])
+		So(merged[1], ShouldEqual, mc1[1])
+		So(merged[2], ShouldEqual, mc2[0])
+		So(merged[3], ShouldEqual, mc2[1])
+	})
+}
+
+func TestServeHTTP(t *testing.T) {
+	t.Parallel()
+	setup := func(h http.Handler) (client *http.Client, ts *httptest.Server) {
+		client = &http.Client{}
+		client.Timeout = 10 * time.Second
+		ts = httptest.NewServer(h)
+		return
+	}
+	cleanup := func(ts *httptest.Server) {
+		ts.Close()
+	}
+	a := func(c *Context, next Handler) {
+		c.Context = context.WithValue(c.Context, "a", "A")
+		next(c)
+	}
+	b := func(c *Context, next Handler) {
+		val := c.Value("a").(string)
+		assert.ShouldEqual(val, "A")
+		c.Context = context.WithValue(c.Context, "a", val+"B")
+		next(c)
+	}
+	handler := func(c *Context) {
+		assert.ShouldEqual(c.Value("b"), "AB")
+		c.WriteHeader(http.StatusOK)
+	}
+	helloHandler := func(c *Context) {
+		c.Write([]byte("Hello, " + c.Params[0].Value))
+	}
+
+	Convey("Middleware registered using Use and Handle should both execute in order", t, func() {
+		r := New()
+		r.Use(MiddlewareChain{a})
+		r.GET("/foo", MiddlewareChain{b}, handler)
+		client, ts := setup(r)
+		defer cleanup(ts)
+
+		res, err := client.Get(ts.URL + "/foo")
+		So(err, ShouldBeNil)
+		res.Body.Close()
+	})
+
+	Convey("Wrong HTTP method for existing path should return method not allowed", t, func() {
+		r := New()
+		r.Use(MiddlewareChain{a})
+		r.GET("/foo", MiddlewareChain{b}, handler)
+		client, ts := setup(r)
+		defer cleanup(ts)
+
+		res, err := client.Post(ts.URL+"/foo", "", nil)
+		So(err, ShouldBeNil)
+		defer res.Body.Close()
+		So(res.StatusCode, ShouldEqual, http.StatusMethodNotAllowed)
+	})
+
+	Convey("Response should match expected value", t, func() {
+		r := New()
+		r.Use(MiddlewareChain{a})
+		r.GET("/hello/:name", MiddlewareChain{b}, helloHandler)
+		client, ts := setup(r)
+		defer cleanup(ts)
+
+		res, err := client.Get(ts.URL + "/hello/世界")
+		So(err, ShouldBeNil)
+		defer res.Body.Close()
+		p, err := ioutil.ReadAll(res.Body)
+		So(res.StatusCode, ShouldEqual, http.StatusOK)
+		So(string(p), ShouldEqual, "Hello, 世界")
+	})
+}
+
+func TestWrite(t *testing.T) {
+	t.Parallel()
+	setup := func(h http.Handler) (client *http.Client, ts *httptest.Server) {
+		client = &http.Client{}
+		client.Timeout = 10 * time.Second
+		ts = httptest.NewServer(h)
+		return
+	}
+	cleanup := func(ts *httptest.Server) {
+		ts.Close()
+	}
+
+	Convey("Written() and StatusCode() should change before and after Write", t, func() {
+		a := func(c *Context, next Handler) {
+			assert.ShouldBeFalse(c.Written())
+			assert.ShouldEqual(c.StatusCode(), 0)
+			next(c)
+			assert.ShouldBeTrue(c.Written())
+			assert.ShouldEqual(c.StatusCode(), http.StatusOK)
+		}
+		b := func(c *Context, next Handler) {
+			assert.ShouldBeFalse(c.Written())
+			assert.ShouldEqual(c.StatusCode(), 0)
+			io.WriteString(c, "foo")
+			assert.ShouldBeTrue(c.Written())
+			assert.ShouldEqual(c.StatusCode(), http.StatusOK)
+			next(c)
+			assert.ShouldBeTrue(c.Written())
+			assert.ShouldEqual(c.StatusCode(), http.StatusOK)
+		}
+		c := func(c *Context, next Handler) {
+			assert.ShouldBeTrue(c.Written())
+			assert.ShouldEqual(c.StatusCode(), http.StatusOK)
+			next(c)
+			assert.ShouldBeTrue(c.Written())
+			assert.ShouldEqual(c.StatusCode(), http.StatusOK)
+		}
+		handler := func(c *Context) {
+			assert.ShouldBeTrue(c.Written())
+			assert.ShouldEqual(c.StatusCode(), http.StatusOK)
+		}
+
+		r := New()
+		r.Use(MiddlewareChain{a, b, c})
+		r.GET("/test/write", nil, handler)
+		client, ts := setup(r)
+		defer cleanup(ts)
+
+		res, err := client.Get(ts.URL + "/test/write")
+		So(err, ShouldBeNil)
+		res.Body.Close()
+	})
+
+	Convey("Written() and StatusCode() should change before and after WriteHeader", t, func() {
+		a := func(c *Context, next Handler) {
+			assert.ShouldBeFalse(c.Written())
+			assert.ShouldEqual(c.StatusCode(), 0)
+			next(c)
+			assert.ShouldBeTrue(c.Written())
+			assert.ShouldEqual(c.StatusCode(), http.StatusTeapot)
+		}
+		b := func(c *Context, next Handler) {
+			assert.ShouldBeFalse(c.Written())
+			assert.ShouldEqual(c.StatusCode(), 0)
+			c.WriteHeader(http.StatusTeapot)
+			assert.ShouldBeTrue(c.Written())
+			assert.ShouldEqual(c.StatusCode(), http.StatusTeapot)
+			next(c)
+			assert.ShouldBeTrue(c.Written())
+			assert.ShouldEqual(c.StatusCode(), http.StatusTeapot)
+		}
+		c := func(c *Context, next Handler) {
+			assert.ShouldBeTrue(c.Written())
+			assert.ShouldEqual(c.StatusCode(), http.StatusTeapot)
+			next(c)
+			assert.ShouldBeTrue(c.Written())
+			assert.ShouldEqual(c.StatusCode(), http.StatusTeapot)
+		}
+		handler := func(c *Context) {
+			assert.ShouldBeTrue(c.Written())
+			assert.ShouldEqual(c.StatusCode(), http.StatusTeapot)
+		}
+
+		r := New()
+		r.Use(MiddlewareChain{a, b, c})
+		r.GET("/test/writeheader", nil, handler)
+		client, ts := setup(r)
+		defer cleanup(ts)
+
+		res, err := client.Get(ts.URL + "/test/writeheader")
+		So(err, ShouldBeNil)
+		res.Body.Close()
+	})
+}