Adds the ability to provide a tripperware to Promtail client. (#3654)

Fixes #3608

Signed-off-by: Cyril Tovena <cyril.tovena@gmail.com>
pull/3658/head
Cyril Tovena 4 years ago committed by GitHub
parent 806d6a503e
commit 04f4f82b8e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
  1. 21
      clients/pkg/promtail/client/client.go
  2. 34
      clients/pkg/promtail/client/client_test.go

@ -157,8 +157,15 @@ type client struct {
cancel context.CancelFunc
}
// Tripperware can wrap a roundtripper.
type Tripperware func(http.RoundTripper) http.RoundTripper
// New makes a new Client.
func New(reg prometheus.Registerer, cfg Config, logger log.Logger) (Client, error) {
return newClient(reg, cfg, logger)
}
func newClient(reg prometheus.Registerer, cfg Config, logger log.Logger) (*client, error) {
if cfg.URL.URL == nil {
return nil, errors.New("client needs target URL")
}
@ -199,6 +206,20 @@ func New(reg prometheus.Registerer, cfg Config, logger log.Logger) (Client, erro
return c, nil
}
// NewWithTripperware creates a new Loki client with a custom tripperware.
func NewWithTripperware(reg prometheus.Registerer, cfg Config, logger log.Logger, tp Tripperware) (Client, error) {
c, err := newClient(reg, cfg, logger)
if err != nil {
return nil, err
}
if tp != nil {
c.client.Transport = tp(c.client.Transport)
}
return c, nil
}
func (c *client) run() {
batches := map[string]*batch{}

@ -1,9 +1,11 @@
package client
import (
"io"
"math"
"net/http"
"net/http/httptest"
"net/url"
"strings"
"testing"
"time"
@ -459,3 +461,35 @@ func createServerHandler(receivedReqsChan chan receivedReq, status int) http.Han
rw.WriteHeader(status)
})
}
type RoundTripperFunc func(*http.Request) (*http.Response, error)
func (r RoundTripperFunc) RoundTrip(req *http.Request) (*http.Response, error) {
return r(req)
}
func Test_Tripperware(t *testing.T) {
url, err := url.Parse("http://foo.com")
require.NoError(t, err)
var called bool
c, err := NewWithTripperware(nil, Config{
URL: flagext.URLValue{URL: url},
}, log.NewNopLogger(), func(rt http.RoundTripper) http.RoundTripper {
return RoundTripperFunc(func(r *http.Request) (*http.Response, error) {
require.Equal(t, r.URL.String(), "http://foo.com")
called = true
return &http.Response{
StatusCode: http.StatusOK,
Body: io.NopCloser(strings.NewReader("ok")),
}, nil
})
})
require.NoError(t, err)
c.Chan() <- api.Entry{
Labels: model.LabelSet{"foo": "bar"},
Entry: logproto.Entry{Timestamp: time.Now(), Line: "foo"},
}
c.Stop()
require.True(t, called)
}

Loading…
Cancel
Save