Add `Header` method to Pusher for custom header (#1218)

Signed-off-by: songjiayang <songjiayang1@gmail.com>
This commit is contained in:
songjiayang 2023-02-07 23:16:37 +08:00 committed by GitHub
parent 2fced96110
commit 2771bcc5da
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 39 additions and 0 deletions

View File

@ -77,6 +77,7 @@ type Pusher struct {
registerer prometheus.Registerer registerer prometheus.Registerer
client HTTPDoer client HTTPDoer
header http.Header
useBasicAuth bool useBasicAuth bool
username, password string username, password string
@ -201,6 +202,13 @@ func (p *Pusher) Client(c HTTPDoer) *Pusher {
return p return p
} }
// Header sets a custom HTTP header for the Pusher's client. For convenience, this method
// returns a pointer to the Pusher itself.
func (p *Pusher) Header(header http.Header) *Pusher {
p.header = header
return p
}
// BasicAuth configures the Pusher to use HTTP Basic Authentication with the // BasicAuth configures the Pusher to use HTTP Basic Authentication with the
// provided username and password. For convenience, this method returns a // provided username and password. For convenience, this method returns a
// pointer to the Pusher itself. // pointer to the Pusher itself.
@ -236,6 +244,9 @@ func (p *Pusher) Delete() error {
if err != nil { if err != nil {
return err return err
} }
if p.header != nil {
req.Header = p.header
}
if p.useBasicAuth { if p.useBasicAuth {
req.SetBasicAuth(p.username, p.password) req.SetBasicAuth(p.username, p.password)
} }
@ -286,6 +297,9 @@ func (p *Pusher) push(ctx context.Context, method string) error {
if err != nil { if err != nil {
return err return err
} }
if p.header != nil {
req.Header = p.header
}
if p.useBasicAuth { if p.useBasicAuth {
req.SetBasicAuth(p.username, p.password) req.SetBasicAuth(p.username, p.password)
} }

View File

@ -31,6 +31,7 @@ func TestPush(t *testing.T) {
lastMethod string lastMethod string
lastBody []byte lastBody []byte
lastPath string lastPath string
lastHeader http.Header
) )
// Fake a Pushgateway that responds with 202 to DELETE and with 200 in // Fake a Pushgateway that responds with 202 to DELETE and with 200 in
@ -38,6 +39,7 @@ func TestPush(t *testing.T) {
pgwOK := httptest.NewServer( pgwOK := httptest.NewServer(
http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
lastMethod = r.Method lastMethod = r.Method
lastHeader = r.Header
var err error var err error
lastBody, err = io.ReadAll(r.Body) lastBody, err = io.ReadAll(r.Body)
if err != nil { if err != nil {
@ -281,4 +283,27 @@ func TestPush(t *testing.T) {
if lastPath != "/metrics/job/testjob/a/x/b/y" && lastPath != "/metrics/job/testjob/b/y/a/x" { if lastPath != "/metrics/job/testjob/a/x/b/y" && lastPath != "/metrics/job/testjob/b/y/a/x" {
t.Error("unexpected path:", lastPath) t.Error("unexpected path:", lastPath)
} }
// Push some Collectors with custom header, all good.
header := make(http.Header)
header.Set("Authorization", "Bearer Token")
if err := New(pgwOK.URL, "testjob").
Collector(metric1).
Collector(metric2).
Header(header).
Push(); err != nil {
t.Fatal(err)
}
if lastMethod != http.MethodPut {
t.Errorf("got method %q for Add, want %q", lastMethod, http.MethodPut)
}
if !bytes.Equal(lastBody, wantBody) {
t.Errorf("got body %v, want %v", lastBody, wantBody)
}
if lastPath != "/metrics/job/testjob" {
t.Error("unexpected path:", lastPath)
}
if lastHeader == nil || lastHeader.Get("Authorization") == "" {
t.Error("empty Authorization header")
}
} }