1
0
mirror of https://github.com/golang/go synced 2024-11-23 20:10:08 -07:00

net/http/httptest: detect Content-Type in ResponseRecorder

* detect Content-Type on ReponseRecorder.Write[String] call
  if header wasn't written yet, Content-Type header is not set and
  Transfer-Encoding is not set.
* fix typos in serve_test.go

Updates #12986

Change-Id: Id2ed8b1994e64657370fed71eb3882d611f76b31
Reviewed-on: https://go-review.googlesource.com/16096
Reviewed-by: Brad Fitzpatrick <bradfitz@golang.org>
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
This commit is contained in:
Nodir Turakulov 2015-10-19 14:36:25 -07:00 committed by Brad Fitzpatrick
parent d4d1668864
commit 0bf515c8c4
3 changed files with 68 additions and 9 deletions

View File

@ -44,11 +44,36 @@ func (rw *ResponseRecorder) Header() http.Header {
return m
}
// writeHeader writes a header if it was not written yet and
// detects Content-Type if needed.
//
// bytes or str are the beginning of the response body.
// We pass both to avoid unnecessarily generate garbage
// in rw.WriteString which was created for performance reasons.
// Non-nil bytes win.
func (rw *ResponseRecorder) writeHeader(b []byte, str string) {
if rw.wroteHeader {
return
}
if len(str) > 512 {
str = str[:512]
}
_, hasType := rw.HeaderMap["Content-Type"]
hasTE := rw.HeaderMap.Get("Transfer-Encoding") != ""
if !hasType && !hasTE {
if b == nil {
b = []byte(str)
}
rw.HeaderMap.Set("Content-Type", http.DetectContentType(b))
}
rw.WriteHeader(200)
}
// Write always succeeds and writes to rw.Body, if not nil.
func (rw *ResponseRecorder) Write(buf []byte) (int, error) {
if !rw.wroteHeader {
rw.WriteHeader(200)
}
rw.writeHeader(buf, "")
if rw.Body != nil {
rw.Body.Write(buf)
}
@ -57,9 +82,7 @@ func (rw *ResponseRecorder) Write(buf []byte) (int, error) {
// WriteString always succeeds and writes to rw.Body, if not nil.
func (rw *ResponseRecorder) WriteString(str string) (int, error) {
if !rw.wroteHeader {
rw.WriteHeader(200)
}
rw.writeHeader(nil, str)
if rw.Body != nil {
rw.Body.WriteString(str)
}
@ -70,8 +93,8 @@ func (rw *ResponseRecorder) WriteString(str string) (int, error) {
func (rw *ResponseRecorder) WriteHeader(code int) {
if !rw.wroteHeader {
rw.Code = code
rw.wroteHeader = true
}
rw.wroteHeader = true
}
// Flush sets rw.Flushed to true.

View File

@ -39,6 +39,14 @@ func TestRecorder(t *testing.T) {
return nil
}
}
hasHeader := func(key, want string) checkFunc {
return func(rec *ResponseRecorder) error {
if got := rec.HeaderMap.Get(key); got != want {
return fmt.Errorf("header %s = %q; want %q", key, got, want)
}
return nil
}
}
tests := []struct {
name string
@ -73,7 +81,12 @@ func TestRecorder(t *testing.T) {
func(w http.ResponseWriter, r *http.Request) {
io.WriteString(w, "hi first")
},
check(hasStatus(200), hasContents("hi first"), hasFlush(false)),
check(
hasStatus(200),
hasContents("hi first"),
hasFlush(false),
hasHeader("Content-Type", "text/plain; charset=utf-8"),
),
},
{
"flush",
@ -83,6 +96,29 @@ func TestRecorder(t *testing.T) {
},
check(hasStatus(200), hasFlush(true)),
},
{
"Content-Type detection",
func(w http.ResponseWriter, r *http.Request) {
io.WriteString(w, "<html>")
},
check(hasHeader("Content-Type", "text/html; charset=utf-8")),
},
{
"no Content-Type detection with Transfer-Encoding",
func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Transfer-Encoding", "some encoding")
io.WriteString(w, "<html>")
},
check(hasHeader("Content-Type", "")), // no header
},
{
"no Content-Type detection if set explicitly",
func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "some/type")
io.WriteString(w, "<html>")
},
check(hasHeader("Content-Type", "some/type")),
},
}
r, _ := http.NewRequest("GET", "http://foo.com/", nil)
for _, tt := range tests {

View File

@ -2487,7 +2487,7 @@ func TestHeaderToWire(t *testing.T) {
if !strings.Contains(got, "404") {
return errors.New("wrong status")
}
if strings.Contains(got, "Some-Header") {
if strings.Contains(got, "Too-Late") {
return errors.New("shouldn't have seen Too-Late")
}
return nil