Files
latticelm/cmd/gateway/main_test.go

58 lines
1.3 KiB
Go

package main
import (
"net/http"
"net/http/httptest"
"testing"
"github.com/stretchr/testify/assert"
)
var _ http.Flusher = (*responseWriter)(nil)
type countingFlusherRecorder struct {
*httptest.ResponseRecorder
flushCount int
}
func newCountingFlusherRecorder() *countingFlusherRecorder {
return &countingFlusherRecorder{ResponseRecorder: httptest.NewRecorder()}
}
func (r *countingFlusherRecorder) Flush() {
r.flushCount++
}
func TestResponseWriterWriteHeaderOnlyOnce(t *testing.T) {
rec := httptest.NewRecorder()
rw := &responseWriter{ResponseWriter: rec, statusCode: http.StatusOK}
rw.WriteHeader(http.StatusCreated)
rw.WriteHeader(http.StatusInternalServerError)
assert.Equal(t, http.StatusCreated, rec.Code)
assert.Equal(t, http.StatusCreated, rw.statusCode)
}
func TestResponseWriterWriteSetsImplicitStatus(t *testing.T) {
rec := httptest.NewRecorder()
rw := &responseWriter{ResponseWriter: rec, statusCode: http.StatusOK}
n, err := rw.Write([]byte("ok"))
assert.NoError(t, err)
assert.Equal(t, 2, n)
assert.Equal(t, http.StatusOK, rec.Code)
assert.Equal(t, http.StatusOK, rw.statusCode)
assert.Equal(t, 2, rw.bytesWritten)
}
func TestResponseWriterFlushDelegates(t *testing.T) {
rec := newCountingFlusherRecorder()
rw := &responseWriter{ResponseWriter: rec, statusCode: http.StatusOK}
rw.Flush()
assert.Equal(t, 1, rec.flushCount)
}