added unit tests for the prometheus middleware

ibetitsmike · ibetitsmike · commit b3b5b75eed6e · 2025-04-15T16:07:30.000Z
diff --git a/coderd/httpmw/prometheus.go b/coderd/httpmw/prometheus.go
@@ -102,6 +102,10 @@ func Prometheus(register prometheus.Registerer) func(http.Handler) http.Handler
 
 func getRoutePattern(r *http.Request) string {
 	rctx := chi.RouteContext(r.Context())
+	if rctx == nil {
+		return ""
+	}
+
 	if pattern := rctx.RoutePattern(); pattern != "" {
 		// Pattern is already available
 		return pattern
@@ -113,7 +117,8 @@ func getRoutePattern(r *http.Request) string {
 	}
 
 	tctx := chi.NewRouteContext()
-	if !rctx.Routes.Match(tctx, r.Method, routePath) {
+	routes := rctx.Routes
+	if routes != nil && routes.Match(tctx, r.Method, routePath) {
 		// No matching pattern, so just return an empty string.
 		// It is done to avoid returning a static path for frontend requests.
 		return ""
diff --git a/coderd/httpmw/prometheus_test.go b/coderd/httpmw/prometheus_test.go
@@ -8,10 +8,14 @@ import (
 
 	"github.com/go-chi/chi/v5"
 	"github.com/prometheus/client_golang/prometheus"
+	cm "github.com/prometheus/client_model/go"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	"github.com/coder/coder/v2/coderd/httpmw"
 	"github.com/coder/coder/v2/coderd/tracing"
+	"github.com/coder/coder/v2/testutil"
+	"github.com/coder/websocket"
 )
 
 func TestPrometheus(t *testing.T) {
@@ -30,3 +34,89 @@ func TestPrometheus(t *testing.T) {
 		require.Greater(t, len(metrics), 0)
 	})
 }
+
+func TestPrometheus_Concurrent(t *testing.T) {
+	t.Parallel()
+	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitShort)
+	defer cancel()
+
+	reg := prometheus.NewRegistry()
+	promMW := httpmw.Prometheus(reg)
+
+	// Create a test handler to simulate a WebSocket connection
+	testHandler := http.HandlerFunc(func(rw http.ResponseWriter, r *http.Request) {
+		conn, err := websocket.Accept(rw, r, nil)
+		if !assert.NoError(t, err, "failed to accept websocket") {
+			return
+		}
+		defer conn.Close(websocket.StatusGoingAway, "")
+	})
+
+	wrappedHandler := promMW(testHandler)
+
+	r := chi.NewRouter()
+	r.Use(tracing.StatusWriterMiddleware, promMW)
+	r.Get("/api/v2/build/{build}/logs", func(rw http.ResponseWriter, r *http.Request) {
+		wrappedHandler.ServeHTTP(rw, r)
+	})
+
+	srv := httptest.NewServer(r)
+	defer srv.Close()
+	// nolint: bodyclose
+	conn, _, err := websocket.Dial(ctx, srv.URL+"/api/v2/build/1/logs", nil)
+	require.NoError(t, err, "failed to dial WebSocket")
+	defer conn.Close(websocket.StatusNormalClosure, "")
+
+	metrics, err := reg.Gather()
+	require.NoError(t, err)
+	require.Greater(t, len(metrics), 0)
+	metricLabels := getMetricLabels(metrics)
+
+	concurrentWebsockets, ok := metricLabels["coderd_api_concurrent_websockets"]
+	require.True(t, ok, "coderd_api_concurrent_websockets metric not found")
+	require.Equal(t, "/api/v2/build/{build}/logs", concurrentWebsockets["path"])
+}
+
+func TestGetRoutePattern_UserRoute(t *testing.T) {
+	t.Parallel()
+	reg := prometheus.NewRegistry()
+	promMW := httpmw.Prometheus(reg)
+
+	r := chi.NewRouter()
+	r.With(promMW).Get("/api/v2/users/{user}", func(w http.ResponseWriter, r *http.Request) {})
+
+	req := httptest.NewRequest("GET", "/api/v2/users/john", nil)
+
+	sw := &tracing.StatusWriter{ResponseWriter: httptest.NewRecorder()}
+
+	r.ServeHTTP(sw, req)
+
+	metrics, err := reg.Gather()
+	require.NoError(t, err)
+	require.Greater(t, len(metrics), 0)
+	metricLabels := getMetricLabels(metrics)
+
+	reqProcessed, ok := metricLabels["coderd_api_requests_processed_total"]
+	require.True(t, ok, "coderd_api_requests_processed_total metric not found")
+	require.Equal(t, "/api/v2/users/{user}", reqProcessed["path"])
+	require.Equal(t, "GET", reqProcessed["method"])
+
+	concurrentRequests, ok := metricLabels["coderd_api_concurrent_requests"]
+	require.True(t, ok, "coderd_api_concurrent_requests metric not found")
+	require.Equal(t, "/api/v2/users/{user}", concurrentRequests["path"])
+	require.Equal(t, "GET", concurrentRequests["method"])
+}
+
+func getMetricLabels(metrics []*cm.MetricFamily) map[string]map[string]string {
+	metricLabels := map[string]map[string]string{}
+	for _, metricFamily := range metrics {
+		metricName := metricFamily.GetName()
+		metricLabels[metricName] = map[string]string{}
+		for _, metric := range metricFamily.GetMetric() {
+			for _, labelPair := range metric.GetLabel() {
+				metricLabels[metricName][labelPair.GetName()] = labelPair.GetValue()
+			}
+		}
+	}
+	return metricLabels
+}
