From e75a244b6127465e7cb39b16513026baf30beca5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1s=20Senart?= <tomas@soundcloud.com>
Date: Fri, 20 Sep 2013 18:17:29 +0200
Subject: [PATCH] Implement Targets#SetHeader

---
 lib/targets.go      |  7 +++++++
 lib/targets_test.go | 11 +++++++++++
 2 files changed, 18 insertions(+)

diff --git a/lib/targets.go b/lib/targets.go
index 9864f86..55fec30 100644
--- a/lib/targets.go
+++ b/lib/targets.go
@@ -56,3 +56,10 @@ func (t Targets) Shuffle(seed int64) {
 		t[i], t[rnd] = t[rnd], t[i]
 	}
 }
+
+// SetHeader sets the passed request header in all Targets
+func (t Targets) SetHeader(header http.Header) {
+	for _, target := range t {
+		target.Header = header
+	}
+}
diff --git a/lib/targets_test.go b/lib/targets_test.go
index 9337928..9212ea3 100644
--- a/lib/targets_test.go
+++ b/lib/targets_test.go
@@ -52,3 +52,14 @@ func TestShuffle(t *testing.T) {
 	}
 	t.Fatal("Targets were not shuffled correctly")
 }
+
+func TestSetHeader(t *testing.T) {
+	targets, _ := NewTargets([]string{"GET http://lolcathost:9999/", "HEAD http://lolcathost:9999/"})
+	want := "lolcathost.com"
+	targets.SetHeader(http.Header{"Host": []string{want}})
+	for _, target := range targets {
+		if got := target.Header.Get("Host"); got != want {
+			t.Errorf("Want: %s, Got: %s", want, got)
+		}
+	}
+}