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 ab35917d492a545cd1da3e44940eea791e2a1f0f..547065f7ba3ef5c54fbb8d814bb4b47d3fe5d104 100644
--- a/appengine/tsmon/middleware_test.go
+++ b/appengine/tsmon/middleware_test.go
@@ -4,55 +4,57 @@
 
 package tsmon
 
 import (
 	"net/http"
 	"net/http/httptest"
 	"testing"
 	"time"
 
 	"github.com/golang/protobuf/proto"
-	"github.com/julienschmidt/httprouter"
 	"github.com/luci/gae/service/datastore"
 	"github.com/luci/luci-go/common/tsmon"
 	"github.com/luci/luci-go/common/tsmon/field"
 	"github.com/luci/luci-go/common/tsmon/store"
 	"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"
-	"golang.org/x/net/context"
+	"github.com/luci/luci-go/server/router"
 
 	. "github.com/smartystreets/goconvey/convey"
 )
 
 func TestMiddleware(t *testing.T) {
 	t.Parallel()
 	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)
-
-		tsmon.Register(c, metric)
-		So(tsmon.Store(c).Incr(c, metric, time.Time{}, []interface{}{}, int64(1)), ShouldBeNil)
+	f := func(c *router.Context) {
+		So(store.IsNilStore(tsmon.Store(c.Context)), ShouldBeFalse)
+		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() {
 		c, _ := buildGAETestContext()
 		state, monitor := buildTestState()
 		ds := datastore.Get(c)
 
 		exists, err := ds.Exists(ds.NewKey("Instance", instanceEntityID(c), 0, nil))
 		So(err, ShouldBeNil)
 		So(exists.All(), ShouldBeFalse)
 
 		rec := httptest.NewRecorder()
-		state.Middleware(f)(c, rec, &http.Request{}, nil)
+		router.RunMiddleware(
+			&router.Context{Context: c, Writer: rec, Request: &http.Request{}},
+			router.MiddlewareChain{state.Middleware},
+			f,
+		)
 		So(rec.Code, ShouldEqual, http.StatusOK)
 
 		exists, err = ds.Exists(ds.NewKey("Instance", instanceEntityID(c), 0, nil))
 		So(err, ShouldBeNil)
 		So(exists.All(), ShouldBeTrue)
 
 		// Shouldn't flush since the instance entity doesn't have a task number yet.
 		So(len(monitor.Cells), ShouldEqual, 0)
 	})
 
@@ -65,21 +67,25 @@ func TestMiddleware(t *testing.T) {
 		i := instance{
 			ID:          instanceEntityID(c),
 			TaskNum:     0,
 			LastUpdated: clock.Now().Add(-2 * time.Minute).UTC(),
 		}
 		So(ds.Put(&i), ShouldBeNil)
 
 		state.lastFlushed = clock.Now().Add(-2 * time.Minute)
 
 		rec := httptest.NewRecorder()
-		state.Middleware(f)(c, rec, &http.Request{}, nil)
+		router.RunMiddleware(
+			&router.Context{Context: c, Writer: rec, Request: &http.Request{}},
+			router.MiddlewareChain{state.Middleware},
+			f,
+		)
 		So(rec.Code, ShouldEqual, http.StatusOK)
 
 		So(len(monitor.Cells), ShouldEqual, 1)
 		So(monitor.Cells[0][0].Name, ShouldEqual, "m")
 		So(monitor.Cells[0][0].Value, ShouldEqual, int64(1))
 
 		// Flushing should update the LastUpdated time.
 		inst, err := getOrCreateInstanceEntity(c)
 		So(err, ShouldBeNil)
 		So(inst.LastUpdated, ShouldResemble, clock.Now().Round(time.Second))
@@ -90,49 +96,60 @@ func TestMiddleware(t *testing.T) {
 		So(value, ShouldEqual, int64(1))
 	})
 
 	Convey("Resets cumulative metrics", t, func() {
 		c, clock := buildGAETestContext()
 		state, monitor := buildTestState()
 
 		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.RunMiddleware(
+			&router.Context{Context: c, Writer: rec, Request: &http.Request{}},
+			router.MiddlewareChain{state.Middleware},
+			func(c *router.Context) {
+				f(c)
+				// 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)
 		So(len(monitor.Cells), ShouldEqual, 0)
 
 		// Value should be reset.
 		value, err := tsmon.Store(c).Get(c, metric, time.Time{}, []interface{}{})
 		So(err, ShouldBeNil)
 		So(value, ShouldBeNil)
 	})
 
 	Convey("Dynamic enable and disable works", t, func() {
 		c, _ := buildGAETestContext()
 		state, _ := buildTestState()
 
 		// 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.RunMiddleware(
+			&router.Context{Context: c, Writer: rec, Request: &http.Request{}},
+			router.MiddlewareChain{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.RunMiddleware(
+			&router.Context{Context: c, Writer: rec, Request: &http.Request{}},
+			router.MiddlewareChain{state.Middleware},
+			func(c *router.Context) {
+				So(store.IsNilStore(tsmon.Store(c.Context)), ShouldBeTrue)
+			},
+		)
 		So(rec.Code, ShouldEqual, http.StatusOK)
 	})
 }