notifier: Allow swapping out HTTP Doer

We need to be able to modify the HTTP POST in Weave Cortex to add multitenancy information to a notification. Since we only really need a special header in the end, the other option would be to just allow passing in headers to the notifier. But swapping out the whole Doer is more general and allows others to swap out the network-talky bits of the notifier for their own use. Doing this via contexts here wouldn't work well, due to the decoupled flow of data in the notifier. There was no existing interface containing the ctxhttp.Post() or ctxhttp.Do() methods, so I settled on just using Do() as a swappable function directly (and with a more minimal signature than Post).
rongfengliang · Feb 27, 2017 · f152ac5 · f152ac5
1 parent 1ab893c
commit f152ac5
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 1 deletion.
diff --git a/notifier/notifier.go b/notifier/notifier.go
@@ -77,12 +77,18 @@ type Options struct {
 	QueueCapacity  int
 	ExternalLabels model.LabelSet
 	RelabelConfigs []*config.RelabelConfig
+	// Used for sending HTTP requests to the Alertmanager.
+	Do func(ctx context.Context, client *http.Client, req *http.Request) (*http.Response, error)
 }
 
 // New constructs a new Notifier.
 func New(o *Options) *Notifier {
 	ctx, cancel := context.WithCancel(context.Background())
 
+	if o.Do == nil {
+		o.Do = ctxhttp.Do
+	}
+
 	return &Notifier{
 		queue:  make(model.Alerts, 0, o.QueueCapacity),
 		ctx:    ctx,
@@ -351,7 +357,12 @@ func (n *Notifier) sendAll(alerts ...*model.Alert) bool {
 }
 
 func (n *Notifier) sendOne(ctx context.Context, c *http.Client, url string, b []byte) error {
-	resp, err := ctxhttp.Post(ctx, c, url, contentTypeJSON, bytes.NewReader(b))
+	req, err := http.NewRequest("POST", url, bytes.NewReader(b))
+	if err != nil {
+		return err
+	}
+	req.Header.Set("Content-Type", contentTypeJSON)
+	resp, err := n.opts.Do(ctx, c, req)
 	if err != nil {
 		return err
 	}

diff --git a/notifier/notifier_test.go b/notifier/notifier_test.go
@@ -16,12 +16,15 @@ package notifier
 import (
 	"encoding/json"
 	"fmt"
+	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
 	"reflect"
 	"testing"
 	"time"
 
+	"golang.org/x/net/context"
+
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/config"
 )
@@ -191,6 +194,37 @@ func TestHandlerSendAll(t *testing.T) {
 	}
 }
 
+func TestCustomDo(t *testing.T) {
+	const testURL = "http://testurl.com/"
+	const testBody = "testbody"
+
+	var received bool
+	h := New(&Options{
+		Do: func(ctx context.Context, client *http.Client, req *http.Request) (*http.Response, error) {
+			received = true
+			body, err := ioutil.ReadAll(req.Body)
+			if err != nil {
+				t.Fatalf("Unable to read request body: %v", err)
+			}
+			if string(body) != testBody {
+				t.Fatalf("Unexpected body; want %v, got %v", testBody, string(body))
+			}
+			if req.URL.String() != testURL {
+				t.Fatalf("Unexpected URL; want %v, got %v", testURL, req.URL.String())
+			}
+			return &http.Response{
+				Body: ioutil.NopCloser(nil),
+			}, nil
+		},
+	})
+
+	h.sendOne(context.Background(), nil, testURL, []byte(testBody))
+
+	if !received {
+		t.Fatal("Expected to receive an alert, but didn't")
+	}
+}
+
 func TestExternalLabels(t *testing.T) {
 	h := New(&Options{
 		QueueCapacity:  3 * maxBatchSize,