diff options
Diffstat (limited to 'workhorse/internal/helper/countingresponsewriter_test.go')
-rw-r--r-- | workhorse/internal/helper/countingresponsewriter_test.go | 50 |
1 files changed, 50 insertions, 0 deletions
diff --git a/workhorse/internal/helper/countingresponsewriter_test.go b/workhorse/internal/helper/countingresponsewriter_test.go new file mode 100644 index 00000000000..f9f2f4ced5b --- /dev/null +++ b/workhorse/internal/helper/countingresponsewriter_test.go @@ -0,0 +1,50 @@ +package helper + +import ( + "bytes" + "io" + "net/http" + "testing" + "testing/iotest" + + "github.com/stretchr/testify/require" +) + +type testResponseWriter struct { + data []byte +} + +func (*testResponseWriter) WriteHeader(int) {} +func (*testResponseWriter) Header() http.Header { return nil } + +func (trw *testResponseWriter) Write(p []byte) (int, error) { + trw.data = append(trw.data, p...) + return len(p), nil +} + +func TestCountingResponseWriterStatus(t *testing.T) { + crw := NewCountingResponseWriter(&testResponseWriter{}) + crw.WriteHeader(123) + crw.WriteHeader(456) + require.Equal(t, 123, crw.Status()) +} + +func TestCountingResponseWriterCount(t *testing.T) { + crw := NewCountingResponseWriter(&testResponseWriter{}) + for _, n := range []int{1, 2, 4, 8, 16, 32} { + _, err := crw.Write(bytes.Repeat([]byte{'.'}, n)) + require.NoError(t, err) + } + require.Equal(t, int64(63), crw.Count()) +} + +func TestCountingResponseWriterWrite(t *testing.T) { + trw := &testResponseWriter{} + crw := NewCountingResponseWriter(trw) + + testData := []byte("test data") + _, err := io.Copy(crw, iotest.OneByteReader(bytes.NewReader(testData))) + require.NoError(t, err) + + require.Equal(t, string(testData), string(trw.data)) +} |