diff --git a/capture_metrics.go b/capture_metrics.go index 97f0a51..9d51185 100644 --- a/capture_metrics.go +++ b/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 } diff --git a/capture_metrics_test.go b/capture_metrics_test.go index 222c8cb..9e76b4a 100644 --- a/capture_metrics_test.go +++ b/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 {