github.com/weaveworks/common@v0.0.0-20230728070032-dd9e68f319d5/middleware/response_test.go (about)

     1  package middleware
     2  
     3  import (
     4  	"bytes"
     5  	"net/http"
     6  	"net/http/httptest"
     7  	"testing"
     8  )
     9  
    10  func TestBadResponseLoggingWriter(t *testing.T) {
    11  	for _, tc := range []struct {
    12  		statusCode int
    13  		data       string
    14  		expected   string
    15  	}{
    16  		{http.StatusOK, "", ""},
    17  		{http.StatusOK, "some data", ""},
    18  		{http.StatusUnprocessableEntity, "unprocessable", ""},
    19  		{http.StatusInternalServerError, "", ""},
    20  		{http.StatusInternalServerError, "bad juju", "bad juju\n"},
    21  	} {
    22  		w := httptest.NewRecorder()
    23  		var buf bytes.Buffer
    24  		wrapped := newBadResponseLoggingWriter(w, &buf)
    25  		switch {
    26  		case tc.data == "":
    27  			wrapped.WriteHeader(tc.statusCode)
    28  		case tc.statusCode < 300 && tc.data != "":
    29  			wrapped.WriteHeader(tc.statusCode)
    30  			wrapped.Write([]byte(tc.data))
    31  		default:
    32  			http.Error(wrapped, tc.data, tc.statusCode)
    33  		}
    34  		if wrapped.getStatusCode() != tc.statusCode {
    35  			t.Errorf("Wrong status code: have %d want %d", wrapped.getStatusCode(), tc.statusCode)
    36  		}
    37  		data := buf.String()
    38  		if data != tc.expected {
    39  			t.Errorf("Wrong data: have %q want %q", data, tc.expected)
    40  		}
    41  	}
    42  }
    43  
    44  // nonFlushingResponseWriter implements http.ResponseWriter but does not implement http.Flusher
    45  type nonFlushingResponseWriter struct{}
    46  
    47  func (rw *nonFlushingResponseWriter) Header() http.Header {
    48  	return nil
    49  }
    50  
    51  func (rw *nonFlushingResponseWriter) Write(_ []byte) (int, error) {
    52  	return -1, nil
    53  }
    54  
    55  func (rw *nonFlushingResponseWriter) WriteHeader(_ int) {
    56  }
    57  
    58  func TestBadResponseLoggingWriter_WithAndWithoutFlusher(t *testing.T) {
    59  	var buf bytes.Buffer
    60  
    61  	nf := newBadResponseLoggingWriter(&nonFlushingResponseWriter{}, &buf)
    62  
    63  	_, ok := nf.(http.Flusher)
    64  	if ok {
    65  		t.Errorf("Should not be able to cast nf as an http.Flusher")
    66  	}
    67  
    68  	rec := httptest.NewRecorder()
    69  	f := newBadResponseLoggingWriter(rec, &buf)
    70  
    71  	ff, ok := f.(http.Flusher)
    72  	if !ok {
    73  		t.Errorf("Should be able to cast f as an http.Flusher")
    74  	}
    75  
    76  	ff.Flush()
    77  	if !rec.Flushed {
    78  		t.Errorf("Flush should have worked but did not")
    79  	}
    80  }
    81  
    82  type responseWriterWithUnwrap interface {
    83  	http.ResponseWriter
    84  	Unwrap() http.ResponseWriter
    85  }
    86  
    87  // Verify that custom http.ResponseWriter implementations implement Unwrap() method, used by http.ResponseContoller.
    88  var _ responseWriterWithUnwrap = &nonFlushingBadResponseLoggingWriter{}
    89  var _ responseWriterWithUnwrap = &flushingBadResponseLoggingWriter{}
    90  var _ responseWriterWithUnwrap = &errorInterceptor{}