8000 Implement Unwrap() for responseWriterDelegator · prometheus/client_golang@9f59f26 · GitHub
[go: up one dir, main page]

Skip to content

Commit 9f59f26

Browse files
igor-drozdovIgor Drozdov
authored andcommitted
Implement Unwrap() for responseWriterDelegator
If the ResponseWriter implements any of the following methods, the ResponseController will call them as appropriate: Flush() FlushError() error // alternative Flush returning an error Hijack() (net.Conn, *bufio.ReadWriter, error) SetReadDeadline(deadline time.Time) error SetWriteDeadline(deadline time.Time) error EnableFullDuplex() error If the ResponseWriter doesn't implement the methods, the ResponseController will call Unwrap() method until it finds a ResponseWriter in the chain This commit implements Unwrap() method to simply return the wrapped ResponseWriter Signed-off-by: Igor Drozdov <ihardrozdov@gmail.com>
1 parent 50ab457 commit 9f59f26

File tree

2 files changed

+96
-0
lines changed

2 files changed

+96
-0
lines changed

prometheus/promhttp/delegator.go

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -76,6 +76,12 @@ func (r *responseWriterDelegator) Write(b []byte) (int, error) {
7676
return n, err
7777
}
7878

79+
// Unwrap lets http.ResponseController get the underlying http.ResponseWriter,
80+
// by implementing the [rwUnwrapper](https://cs.opensource.google/go/go/+/refs/tags/go1.21.4:src/net/http/responsecontroller.go;l=42-44) interface.
81+
func (r *responseWriterDelegator) Unwrap() http.ResponseWriter {
82+
return r.ResponseWriter
83+
}
84+
7985
type (
8086
closeNotifierDelegator struct{ *responseWriterDelegator }
8187
flusherDelegator struct{ *responseWriterDelegator }
Lines changed: 90 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,90 @@
1+
// Copyright 2024 The Prometheus Authors
2+
// Licensed under the Apache License, Version 2.0 (the "License");
3+
// you may not use this file except in compliance with the License.
4+
// You may obtain a copy of the License at
5+
//
6+
// http://www.apache.org/licenses/LICENSE-2.0
7+
//
8+
// Unless required by applicable law or agreed to in writing, software
9+
// distributed under the License is distributed on an "AS IS" BASIS,
10+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
// See the License for the specific language governing permissions and
12+
// limitations under the License.
13+
14+
package promhttp
15+
16+
import (
17+
"net/http"
18+
"time"
19+
20+
"testing"
21+
)
22+
23+
type responseWriter struct {
24+
flushErrorCalled bool
25+
enableFullDuplexCalled bool
26+
setWriteDeadlineCalled time.Time
27+
setReadDeadlineCalled time.Time
28+
}
29+
30+
func (rw *responseWriter) Header() http.Header {
31+
return nil
32+
}
33+
34+
func (rw *responseWriter) Write(p []byte) (int, error) {
35+
return 0, nil
36+
}
37+
38+
func (rw *responseWriter) WriteHeader(statusCode int) {
39+
}
40+
41+
func (rw *responseWriter) FlushError() error {
42+
rw.flushErrorCalled = true
43+
44+
return nil
45+
}
46+
47+
func (rw *responseWriter) SetWriteDeadline(deadline time.Time) error {
48+
rw.setWriteDeadlineCalled = deadline
49+
50+
return nil
51+
}
52+
53+
func (rw *responseWriter) SetReadDeadline(deadline time.Time) error {
54+
rw.setReadDeadlineCalled = deadline
55+
56+
return nil
57+
}
58+
59+
func (rw *responseWriter) EnableFullDuplex() error {
60+
rw.enableFullDuplexCalled = true
61+
62+
return nil
63+
}
64+
65+
func TestResponseWriterDelegatorUnwrap(t *testing.T) {
66+
w := &responseWriter{}
67+
rwd := &responseWriterDelegator{ResponseWriter: w}
68+
69+
if rwd.Unwrap() != w {
70+
t.Error("unwrapped responsewriter must equal to the original responsewriter")
71+
}
72+
73+
controller := http.NewResponseController(rwd)
74+
if err := controller.Flush(); err != nil || !w.flushErrorCalled {
75+
t.Error("FlushError must be propagated to the original responsewriter")
76+
}
77+
78+
if err := controller.EnableFullDuplex(); err != nil || !w.enableFullDuplexCalled {
79+
t.Error("EnableFullDuplex must be propagated to the original responsewriter")
80+
}
81+
82+
timeNow := time.Now()
83+
if err := controller.SetWriteDeadline(timeNow); err != nil || w.setWriteDeadlineCalled != timeNow {
84+
t.Error("SetWriteDeadline must be propagated to the original responsewriter")
85+
}
86+
87+
if err := controller.SetReadDeadline(timeNow); err != nil || w.setReadDeadlineCalled != timeNow {
88+
t.Error("SetReadDeadline must be propagated to the original responsewriter")
89+
}
90+
}

0 commit comments

Comments
 (0)
0