66 lines
1.6 KiB
Go
66 lines
1.6 KiB
Go
package observability
|
|
|
|
import (
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
)
|
|
|
|
var _ http.Flusher = (*metricsResponseWriter)(nil)
|
|
var _ http.Flusher = (*statusResponseWriter)(nil)
|
|
|
|
type testFlusherRecorder struct {
|
|
*httptest.ResponseRecorder
|
|
flushCount int
|
|
}
|
|
|
|
func newTestFlusherRecorder() *testFlusherRecorder {
|
|
return &testFlusherRecorder{ResponseRecorder: httptest.NewRecorder()}
|
|
}
|
|
|
|
func (r *testFlusherRecorder) Flush() {
|
|
r.flushCount++
|
|
}
|
|
|
|
func TestMetricsResponseWriterWriteHeaderOnlyOnce(t *testing.T) {
|
|
rec := httptest.NewRecorder()
|
|
rw := &metricsResponseWriter{ResponseWriter: rec, statusCode: http.StatusOK}
|
|
|
|
rw.WriteHeader(http.StatusAccepted)
|
|
rw.WriteHeader(http.StatusInternalServerError)
|
|
|
|
assert.Equal(t, http.StatusAccepted, rec.Code)
|
|
assert.Equal(t, http.StatusAccepted, rw.statusCode)
|
|
}
|
|
|
|
func TestMetricsResponseWriterFlushDelegates(t *testing.T) {
|
|
rec := newTestFlusherRecorder()
|
|
rw := &metricsResponseWriter{ResponseWriter: rec, statusCode: http.StatusOK}
|
|
|
|
rw.Flush()
|
|
|
|
assert.Equal(t, 1, rec.flushCount)
|
|
}
|
|
|
|
func TestStatusResponseWriterWriteHeaderOnlyOnce(t *testing.T) {
|
|
rec := httptest.NewRecorder()
|
|
rw := &statusResponseWriter{ResponseWriter: rec, statusCode: http.StatusOK}
|
|
|
|
rw.WriteHeader(http.StatusNoContent)
|
|
rw.WriteHeader(http.StatusInternalServerError)
|
|
|
|
assert.Equal(t, http.StatusNoContent, rec.Code)
|
|
assert.Equal(t, http.StatusNoContent, rw.statusCode)
|
|
}
|
|
|
|
func TestStatusResponseWriterFlushDelegates(t *testing.T) {
|
|
rec := newTestFlusherRecorder()
|
|
rw := &statusResponseWriter{ResponseWriter: rec, statusCode: http.StatusOK}
|
|
|
|
rw.Flush()
|
|
|
|
assert.Equal(t, 1, rec.flushCount)
|
|
}
|