Return instrumented response writer from context

This is ensures that users of the ResponseWriter from the context correctly
track usage. Otherwise, context reporting is incorrect.

Signed-off-by: Stephen J Day <stephen.day@docker.com>
This commit is contained in:
Stephen J Day 2015-04-20 11:15:01 -07:00
parent a9a9658e41
commit cb938c476f
2 changed files with 12 additions and 3 deletions

View file

@ -302,7 +302,7 @@ func (irw *instrumentedResponseWriter) Flush() {
func (irw *instrumentedResponseWriter) Value(key interface{}) interface{} { func (irw *instrumentedResponseWriter) Value(key interface{}) interface{} {
if keyStr, ok := key.(string); ok { if keyStr, ok := key.(string); ok {
if keyStr == "http.response" { if keyStr == "http.response" {
return irw.ResponseWriter return irw
} }
if !strings.HasPrefix(keyStr, "http.response.") { if !strings.HasPrefix(keyStr, "http.response.") {

View file

@ -132,8 +132,17 @@ func TestWithResponseWriter(t *testing.T) {
trw := testResponseWriter{} trw := testResponseWriter{}
ctx, rw := WithResponseWriter(Background(), &trw) ctx, rw := WithResponseWriter(Background(), &trw)
if ctx.Value("http.response") != &trw { if ctx.Value("http.response") != rw {
t.Fatalf("response not available in context: %v != %v", ctx.Value("http.response"), &trw) t.Fatalf("response not available in context: %v != %v", ctx.Value("http.response"), rw)
}
grw, err := GetResponseWriter(ctx)
if err != nil {
t.Fatalf("error getting response writer: %v", err)
}
if grw != rw {
t.Fatalf("unexpected response writer returned: %#v != %#v", grw, rw)
} }
if n, err := rw.Write(make([]byte, 1024)); err != nil { if n, err := rw.Write(make([]byte, 1024)); err != nil {