From 3bc8e4c00b65d88d708e9f5cb87d7425486583ae Mon Sep 17 00:00:00 2001
From: Linas Naginionis <linas.n@chronosphere.io>
Date: Mon, 16 Nov 2020 14:42:09 +0200
Subject: [PATCH] More handler tests.

---
 src/query/api/v1/httpd/handler_test.go | 58 +++++++++++++++++++++-----
 1 file changed, 47 insertions(+), 11 deletions(-)

diff --git a/src/query/api/v1/httpd/handler_test.go b/src/query/api/v1/httpd/handler_test.go
index 148abe7b29..69ad5ba9ae 100644
--- a/src/query/api/v1/httpd/handler_test.go
+++ b/src/query/api/v1/httpd/handler_test.go
@@ -377,22 +377,24 @@ func init() {
 	testWorkerPool.Init()
 }
 
-type assertFn func(t *testing.T, prev http.Handler)
+type assertFn func(t *testing.T, prev http.Handler, r *http.Request)
 
 type customHandler struct {
 	t *testing.T
+	routeName string
+	methods []string
 	assertFn assertFn
 }
 
-func (h *customHandler) Route() string     { return "/custom" }
-func (h *customHandler) Methods() []string { return []string{http.MethodGet} }
+func (h *customHandler) Route() string     { return h.routeName }
+func (h *customHandler) Methods() []string { return h.methods }
 func (h *customHandler) Handler(
 	opts options.HandlerOptions,
 	prev http.Handler,
 ) (http.Handler, error) {
 	assert.Equal(h.t, "z", string(opts.TagOptions().MetricName()))
 	fn := func(w http.ResponseWriter, r *http.Request) {
-		h.assertFn(h.t, prev)
+		h.assertFn(h.t, prev, r)
 		_, err := w.Write([]byte("success!"))
 		require.NoError(h.t, err)
 	}
@@ -401,8 +403,6 @@ func (h *customHandler) Handler(
 }
 
 func TestCustomRoutes(t *testing.T) {
-	req := httptest.NewRequest(http.MethodGet, "/custom", nil)
-	res := httptest.NewRecorder()
 	ctrl := gomock.NewController(t)
 	store, _ := m3.NewStorageAndSession(t, ctrl)
 	instrumentOpts := instrument.NewOptions()
@@ -419,22 +419,58 @@ func TestCustomRoutes(t *testing.T) {
 	require.NoError(t, err)
 	custom := &customHandler{
 		t: t,
-		assertFn: func(t *testing.T, prev http.Handler) {
+		routeName: "/custom",
+		methods: []string{http.MethodGet, http.MethodHead},
+		assertFn: func(t *testing.T, prev http.Handler, r *http.Request) {
 			assert.Nil(t, prev, "Should not shadow already existing handler")
 		},
 	}
-	custom2 := &customHandler{
+	customShadowGet := &customHandler{
+		t: t,
+		routeName: "/custom",
+		methods: []string{http.MethodGet},
+		assertFn: func(t *testing.T, prev http.Handler, r *http.Request) {
+			assert.NotNil(t, prev, "Should shadow already existing handler")
+		},
+	}
+	customShadowHead := &customHandler{
 		t: t,
-		assertFn: func(t *testing.T, prev http.Handler) {
+		routeName: "/custom",
+		methods: []string{http.MethodHead},
+		assertFn: func(t *testing.T, prev http.Handler, r *http.Request) {
 			assert.NotNil(t, prev, "Should shadow already existing handler")
 		},
 	}
-	handler := NewHandler(opts, custom, custom2)
+	customNew := &customHandler{
+		t: t,
+		routeName: "/custom/new",
+		methods: []string{http.MethodGet, http.MethodHead},
+		assertFn: func(t *testing.T, prev http.Handler, r *http.Request) {
+			assert.Nil(t, prev, "Should not shadow already existing handler")
+		},
+	}
+	handler := NewHandler(opts, custom, customShadowGet, customShadowHead, customNew)
 	require.NoError(t, err, "unable to setup handler")
 	err = handler.RegisterRoutes()
 	require.NoError(t, err, "unable to register routes")
+
+	for _, method := range custom.methods {
+		assertRoute(t, custom.routeName, method, handler, http.StatusOK)
+	}
+
+	for _, method := range customNew.methods {
+		assertRoute(t, customNew.routeName, method, handler, http.StatusOK)
+	}
+
+	assertRoute(t, customNew.routeName, http.MethodPost, handler, http.StatusMethodNotAllowed)
+	assertRoute(t, "/unknown", http.MethodGet, handler, http.StatusNotFound)
+}
+
+func assertRoute(t *testing.T, routeName string, method string, handler *Handler, expectedStatusCode int) {
+	req := httptest.NewRequest(method, routeName, nil)
+	res := httptest.NewRecorder()
 	handler.Router().ServeHTTP(res, req)
-	require.Equal(t, res.Code, http.StatusOK)
+	require.Equal(t, expectedStatusCode, res.Code)
 }
 
 func TestRouteName(t *testing.T) {