Record HTTP 101 in captured metrics

capture_metrics.go

@@ -8,9 +8,9 @@ import (
 
 // Metrics holds metrics captured from CaptureMetrics.
 type Metrics struct {
-	// Code is the first http response code passed to the WriteHeader func of
-	// the ResponseWriter. If no such call is made, a default code of 200 is
-	// assumed instead.
+	// Code is the first final HTTP response code passed to WriteHeader,
+	// including 101 Switching Protocols. If no final response is written,
+	// a default code of 200 is assumed instead.
 	Code int
 	// Duration is the time it took to execute the handler.
 	Duration time.Duration
@@ -52,7 +52,7 @@ func (m *Metrics) CaptureMetrics(w http.ResponseWriter, fn func(http.ResponseWri
 				return func(code int) {
 					next(code)
 
-					if !(code >= 100 && code <= 199) && !headerWritten {
+					if !headerWritten && (code < 100 || code > 199 || code == http.StatusSwitchingProtocols) {
 						m.Code = code
 						headerWritten = true
 					}

capture_metrics_test.go

@@ -31,6 +31,56 @@ func TestCaptureMetrics(t *testing.T) {
 			Handler:  http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}),
 			WantCode: http.StatusOK,
 		},
+		{
+			Name: "switching protocols",
+			Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(http.StatusSwitchingProtocols)
+			}),
+			WantCode: http.StatusSwitchingProtocols,
+		},
+		{
+			Name: "early hints before switching protocols",
+			Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(http.StatusEarlyHints)
+				w.WriteHeader(http.StatusSwitchingProtocols)
+			}),
+			WantCode: http.StatusSwitchingProtocols,
+		},
+		{
+			Name: "multiple informational headers before switching protocols",
+			Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(http.StatusContinue)
+				w.WriteHeader(http.StatusEarlyHints)
+				w.WriteHeader(http.StatusSwitchingProtocols)
+			}),
+			WantCode: http.StatusSwitchingProtocols,
+		},
+		{
+			Name: "header after switching protocols",
+			Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(http.StatusSwitchingProtocols)
+				w.WriteHeader(http.StatusNoContent)
+			}),
+			WantCode: http.StatusSwitchingProtocols,
+		},
+		{
+			Name: "early hints before final header",
+			Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(http.StatusEarlyHints)
+				w.WriteHeader(http.StatusNoContent)
+			}),
+			WantCode: http.StatusNoContent,
+		},
+		{
+			Name: "informational headers before body",
+			Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(http.StatusContinue)
+				w.WriteHeader(http.StatusEarlyHints)
+				w.Write([]byte("foo"))
+			}),
+			WantCode:    http.StatusOK,
+			WantWritten: 3,
+		},
 		{
 			Name: "headers and body",
 			Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -107,11 +157,14 @@ func TestCaptureMetrics(t *testing.T) {
 			s := httptest.NewServer(h)
 			defer s.Close()
 			res, err := http.Get(s.URL)
-			if !errContains(err, test.WantErr) {
+			if (test.WantErr == "" && err != nil) || !errContains(err, test.WantErr) {
 				t.Errorf("test %d: got=%s want=%s", i, err, test.WantErr)
 			}
 			if err == nil {
 				defer res.Body.Close()
+				if res.StatusCode != test.WantCode {
+					t.Errorf("test %d: response code=%d want=%d", i, res.StatusCode, test.WantCode)
+				}
 			}
 			m := <-ch
 			if m.Code != test.WantCode {