Make HTTP middleware easier to use

appengine/tsmon/middleware_test.go

Index: appengine/tsmon/middleware_test.go
diff --git a/appengine/tsmon/middleware_test.go b/appengine/tsmon/middleware_test.go
index 986aaabdb0a9616308b90d3048684279473d159c..65973f7173e5836fc0585da0df8c35d9f1205944 100644
--- a/appengine/tsmon/middleware_test.go
+++ b/appengine/tsmon/middleware_test.go
@@ -19,6 +19,7 @@ import (
 	"github.com/luci/luci-go/common/tsmon/store/storetest"
 	"github.com/luci/luci-go/common/tsmon/target"
 	"github.com/luci/luci-go/common/tsmon/types"
+	"github.com/luci/luci-go/server/router"
 	"golang.org/x/net/context"
 
 	. "github.com/smartystreets/goconvey/convey"
@@ -27,11 +28,20 @@ import (
 func TestMiddleware(t *testing.T) {
 	metric := &storetest.FakeMetric{"m", "", []field.Field{}, types.CumulativeIntType}
 
-	f := func(c context.Context, rw http.ResponseWriter, r *http.Request, p httprouter.Params) {
-		So(store.IsNilStore(tsmon.Store(c)), ShouldBeFalse)
+	initialize := func(c context.Context, w http.ResponseWriter, r *http.Request, p httprouter.Params) router.Handler {
+		return func(ctx *router.Context) {
+			ctx.Context = c
+			ctx.Writer = w
+			ctx.Request = r
+			ctx.Params = p
+		}
+	}
+
+	f := func(c *router.Context) {
+		So(store.IsNilStore(tsmon.Store(c.Context)), ShouldBeFalse)
 
-		tsmon.Register(c, metric)
-		So(tsmon.Store(c).Incr(c, metric, time.Time{}, []interface{}{}, int64(1)), ShouldBeNil)
+		tsmon.Register(c.Context, metric)
+		So(tsmon.Store(c.Context).Incr(c.Context, metric, time.Time{}, []interface{}{}, int64(1)), ShouldBeNil)
 	}
 
 	Convey("Creates instance entity", t, func() {
@@ -44,7 +54,7 @@ func TestMiddleware(t *testing.T) {
 		So(exists.All(), ShouldBeFalse)
 
 		rec := httptest.NewRecorder()
-		state.Middleware(f)(c, rec, &http.Request{}, nil)
+		router.ChainHandlers(initialize(c, rec, &http.Request{}, nil), state.Middleware(), f)()
 		So(rec.Code, ShouldEqual, http.StatusOK)
 
 		exists, err = ds.Exists(ds.NewKey("Instance", instanceEntityID(c), 0, nil))
@@ -71,7 +81,7 @@ func TestMiddleware(t *testing.T) {
 		state.lastFlushed = clock.Now().Add(-2 * time.Minute)
 
 		rec := httptest.NewRecorder()
-		state.Middleware(f)(c, rec, &http.Request{}, nil)
+		router.ChainHandlers(initialize(c, rec, &http.Request{}, nil), state.Middleware(), f)()
 		So(rec.Code, ShouldEqual, http.StatusOK)
 
 		So(len(monitor.Cells), ShouldEqual, 1)
@@ -96,15 +106,18 @@ func TestMiddleware(t *testing.T) {
 		state.lastFlushed = clock.Now().Add(-2 * time.Minute)
 
 		rec := httptest.NewRecorder()
-		state.Middleware(func(c context.Context, rw http.ResponseWriter, r *http.Request, p httprouter.Params) {
-			f(c, rw, r, p)
-
-			// Override the TaskNum here - it's created just before this handler runs
-			// and used just after.
-			tar := tsmon.Store(c).DefaultTarget().(*target.Task)
-			tar.TaskNum = proto.Int32(int32(0))
-			tsmon.Store(c).SetDefaultTarget(tar)
-		})(c, rec, &http.Request{}, nil)
+		router.ChainHandlers(
+			initialize(c, rec, &http.Request{}, nil),
+			state.Middleware(),
+			f,
+			func(c *router.Context) {
+				// Override the TaskNum here - it's created just before this handler runs
+				// and used just after.
+				tar := tsmon.Store(c.Context).DefaultTarget().(*target.Task)
+				tar.TaskNum = proto.Int32(int32(0))
+				tsmon.Store(c.Context).SetDefaultTarget(tar)
+			},
+		)()
 		So(rec.Code, ShouldEqual, http.StatusOK)
 
 		So(len(tsmon.GetState(c).RegisteredMetrics), ShouldEqual, 1)
@@ -122,16 +135,24 @@ func TestMiddleware(t *testing.T) {
 
 		// Enabled. Store is not nil.
 		rec := httptest.NewRecorder()
-		state.Middleware(func(c context.Context, rw http.ResponseWriter, r *http.Request, p httprouter.Params) {
-			So(store.IsNilStore(tsmon.Store(c)), ShouldBeFalse)
-		})(c, rec, &http.Request{}, nil)
+		router.ChainHandlers(
+			initialize(c, rec, &http.Request{}, nil),
+			state.Middleware(),
+			func(c *router.Context) {
+				So(store.IsNilStore(tsmon.Store(c.Context)), ShouldBeFalse)
+			},
+		)()
 		So(rec.Code, ShouldEqual, http.StatusOK)
 
 		// Disabled. Store is nil.
 		state.testingSettings.Enabled = false
-		state.Middleware(func(c context.Context, rw http.ResponseWriter, r *http.Request, p httprouter.Params) {
-			So(store.IsNilStore(tsmon.Store(c)), ShouldBeTrue)
-		})(c, rec, &http.Request{}, nil)
+		router.ChainHandlers(
+			initialize(c, rec, &http.Request{}, nil),
+			state.Middleware(),
+			func(c *router.Context) {
+				So(store.IsNilStore(tsmon.Store(c.Context)), ShouldBeTrue)
+			},
+		)()
 		So(rec.Code, ShouldEqual, http.StatusOK)
 	})
 }