graceful test no longer sleeps, uses polling

logger/logger.go

@@ -109,6 +109,7 @@ func Middleware(logger *slog.Logger, shouldPrint bool) func(next http.Handler) h
 
 			start := time.Now()
 			metrics.InFlightGauge.Inc()
+			logger.Error("should increment flight gauge")
 
 			defer func() {
 				metrics.InFlightGauge.Dec()
@@ -124,14 +125,17 @@ func Middleware(logger *slog.Logger, shouldPrint bool) func(next http.Handler) h
 
 				metrics.RequestCount.With(prometheus.Labels{"method": r.Method, "endpoint": path}).Inc()
 				metrics.RequestDuration.With(prometheus.Labels{"method": r.Method, "endpoint": path}).Observe(duration.Seconds())
-				logger.Info("route",
-					"path", path,
-					"verb", r.Method,
-					"status", http.StatusText(ww.Status()),
-					"code", ww.Status(),
-					"bytes", ww.BytesWritten(),
-					"duration_ms", duration.Milliseconds(),
-				)
+
+				if shouldPrint {
+					logger.Info("route",
+						"path", path,
+						"verb", r.Method,
+						"status", http.StatusText(ww.Status()),
+						"code", ww.Status(),
+						"bytes", ww.BytesWritten(),
+						"duration_ms", duration.Milliseconds(),
+					)
+				}
 			}()
 
 			next.ServeHTTP(ww, r)

server/server.go

@@ -223,8 +223,7 @@ func customCORSMiddleware(allowedOrigins []string) func(http.Handler) http.Handl
 			w.Header().Set("Access-Control-Allow-Methods", "GET, POST, PUT, DELETE, OPTIONS")
 			w.Header().Set("Access-Control-Allow-Headers", strings.Join([]string{
 				"Accept", "Authorization", "Content-Type", "X-CSRF-Token",
-				"Accept-Language", "Hx-Current-Url", "Hx-Request", "Hx-Target",
-				"Hx-Trigger", "Referer", "User-Agent",
+				"Accept-Language", "Referer", "User-Agent",
 			}, ","))
 			w.Header().Set("Access-Control-Expose-Headers", "Link")
 			w.Header().Set("Access-Control-Max-Age", "300")
@@ -255,7 +254,7 @@ func (s *Server) Serve() error {
 	})
 
 	// all application routes should be defined below
-	router := s.newRouter().With(s.loggerMiddleware)
+	router := s.newRouter()
 
 	// if routes require authentication, use a new With or add it above as a separate middleware
 	// router.Get("/", s.uiIndex)

server/server_test.go

@@ -1,11 +1,14 @@
 package server
 
 import (
+	"bufio"
 	"bytes"
 	"fmt"
 	"io"
 	"log/slog"
 	"net/http"
+	"strconv"
+	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -48,8 +51,10 @@ func TestGracefulShutdown(t *testing.T) {
 	srv, err := newTestServer()
 	require.NoError(t, err)
 
+	concurrentRequests := 10
+
 	wg := sync.WaitGroup{}
-	for i := 0; i < 3; i++ {
+	for i := 0; i < concurrentRequests; i++ {
 		wg.Add(1)
 		go func() {
 			source := fmt.Sprintf("http://localhost:%d/?delay=1s", srv.Port())
@@ -61,8 +66,10 @@ func TestGracefulShutdown(t *testing.T) {
 	}
 
 	// give time for all requests to go out
-	// todo: instead of sleeping in tests, use the metrics endpoint to show inflight request
-	time.Sleep(100 * time.Millisecond)
+	// note: avoid sleeping in tests. To keep tests fast, poll for the expected state or event over a channel
+	// AVOID: time.Sleep(1000 * time.Millisecond)
+	// ATTEMPT: poll for the expected state
+	assertMetric(t, srv, "http_in_flight_requests", float64(concurrentRequests), 2*time.Second)
 
 	err = srv.Close()
 	require.NoError(t, err)
@@ -75,6 +82,74 @@ func TestGracefulShutdown(t *testing.T) {
 	wg.Wait()
 }
 
+func assertMetric(t *testing.T, srv *Server, metric string, target float64, timeout time.Duration) {
+	start := time.Now()
+	for {
+		value, err := getMetric(srv, metric)
+		if err == nil && value == target {
+			return
+		} else if err != nil {
+			t.Errorf("error fetching metric: %v", err)
+			return
+		}
+
+		fmt.Printf("waiting for metric %s to reach %f, currently at %f\n", metric, target, value)
+
+		// Check if timeout has been reached
+		if time.Since(start) >= timeout {
+			t.Errorf("timeout reached before target metric value was reached: %s=%f, got %s=%f", metric, target, metric, value)
+			return
+		}
+
+		// Wait for the next interval
+		time.Sleep(100 * time.Millisecond)
+	}
+}
+
+func getMetric(srv *Server, metric string) (float64, error) {
+	resp, err := http.Get(fmt.Sprintf("http://localhost:%d/metrics", srv.InternalPort()))
+	if err != nil {
+		return 0, err
+	}
+
+	buf := new(bytes.Buffer)
+	buf.ReadFrom(resp.Body)
+	resp.Body.Close()
+
+	return findMetricValue(buf, metric)
+}
+
+// findMetricValue parses the metrics data and retrieves the number from the first line with the given prefix.
+func findMetricValue(metrics *bytes.Buffer, prefix string) (float64, error) {
+	scanner := bufio.NewScanner(metrics)
+	for scanner.Scan() {
+		line := scanner.Text()
+		// Skip comment lines
+		if strings.HasPrefix(line, "#") {
+			continue
+		}
+		// Check if the line starts with the desired prefix
+		if strings.HasPrefix(line, prefix) {
+			// Split the line into the metric name and value
+			parts := strings.Fields(line)
+			if len(parts) < 2 {
+				return 0, fmt.Errorf("malformed metric line: %s", line)
+			}
+			// Convert the value to a float
+			value, err := strconv.ParseFloat(parts[1], 64)
+			if err != nil {
+				return 0, fmt.Errorf("invalid metric value: %v", err)
+			}
+			return value, nil
+		}
+	}
+	// Return an error if no matching prefix was found
+	if err := scanner.Err(); err != nil {
+		return 0, fmt.Errorf("error reading metrics: %v", err)
+	}
+	return 0, fmt.Errorf("metric with prefix '%s' not found", prefix)
+}
+
 // newTestServer is generally called with no parameter. A bit of a hack on variadics, but if you want to pass in a buffer, pass one in.
 // var buf bytes.Buffer
 // newTestServer(buf)