github.com/rudderlabs/rudder-go-kit@v0.30.0/chiware/stats_test.go (about)

     1  package chiware_test
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"net/http"
     7  	"net/http/httptest"
     8  	"strconv"
     9  	"testing"
    10  
    11  	"github.com/go-chi/chi/v5"
    12  	"github.com/golang/mock/gomock"
    13  	"github.com/stretchr/testify/require"
    14  
    15  	"github.com/rudderlabs/rudder-go-kit/chiware"
    16  	"github.com/rudderlabs/rudder-go-kit/stats"
    17  	"github.com/rudderlabs/rudder-go-kit/stats/mock_stats"
    18  )
    19  
    20  func TestStatsMiddleware(t *testing.T) {
    21  	component := "test"
    22  	testCase := func(expectedStatusCode int, pathTemplate, requestPath, expectedMethod string, options ...chiware.Option) func(t *testing.T) {
    23  		return func(t *testing.T) {
    24  			ctrl := gomock.NewController(t)
    25  			mockStats := mock_stats.NewMockStats(ctrl)
    26  			handler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    27  				w.WriteHeader(expectedStatusCode)
    28  			})
    29  
    30  			measurement := mock_stats.NewMockMeasurement(ctrl)
    31  			mockStats.EXPECT().NewStat(fmt.Sprintf("%s.concurrent_requests_count", component), stats.GaugeType).Return(measurement).Times(1)
    32  			mockStats.EXPECT().NewSampledTaggedStat(fmt.Sprintf("%s.response_time", component), stats.TimerType,
    33  				map[string]string{
    34  					"reqType": pathTemplate,
    35  					"method":  expectedMethod,
    36  					"code":    strconv.Itoa(expectedStatusCode),
    37  				}).Return(measurement).Times(1)
    38  			measurement.EXPECT().Since(gomock.Any()).Times(1)
    39  
    40  			ctx, cancel := context.WithCancel(context.Background())
    41  			defer cancel()
    42  			router := chi.NewRouter()
    43  			router.Use(
    44  				chiware.StatMiddleware(ctx, mockStats, component, options...),
    45  			)
    46  			router.MethodFunc(expectedMethod, pathTemplate, handler)
    47  
    48  			response := httptest.NewRecorder()
    49  			request := httptest.NewRequest("GET", "http://example.com"+requestPath, http.NoBody)
    50  			router.ServeHTTP(response, request)
    51  			require.Equal(t, expectedStatusCode, response.Code)
    52  		}
    53  	}
    54  
    55  	t.Run("template with param in path", testCase(http.StatusNotFound, "/v1/{param}", "/v1/abc", "GET"))
    56  	t.Run("template without param in path", testCase(http.StatusNotFound, "/v1/some-other/key", "/v1/some-other/key", "GET"))
    57  	t.Run("template with unknown path ", testCase(http.StatusNotFound, "/a/b/c", "/a/b/c", "GET", chiware.RedactUnknownPaths(false)))
    58  	t.Run("template with unknown path ", testCase(http.StatusNotFound, "/redacted", "/a/b/c", "GET", chiware.RedactUnknownPaths(true)))
    59  	t.Run("template with unknown path ", testCase(http.StatusNotFound, "/redacted", "/a/b/c", "GET"))
    60  }