Files
latticelm/internal/observability/middleware_response_writer_test.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)
}