Make HTTP middleware easier to use

server/router/router_test.go

+// 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()
+        })
+}