From f152ac5e23943fd22243896533df626ba89f3d54 Mon Sep 17 00:00:00 2001
From: Julius Volz <julius@soundcloud.com>
Date: Mon, 27 Feb 2017 20:31:16 +0100
Subject: [PATCH] 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).
---
 notifier/notifier.go      | 13 ++++++++++++-
 notifier/notifier_test.go | 34 ++++++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/notifier/notifier.go b/notifier/notifier.go
index b896d53a40..e9af00a4ca 100644
--- 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
index 65c49a3fad..3587a1a8a7 100644
--- 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,