From c1e496297227da1465cc794ca46d70988dda9eea Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1s=20Senart?= <tomas@soundcloud.com>
Date: Wed, 14 Aug 2013 17:30:13 +0200
Subject: [PATCH] Extract main logic into attack function and test it

---
 main.go      | 62 +++++++++++++++++++++++++++-------------------------
 main_test.go | 27 +++++++++++++++++++++++
 reporters.go |  4 ++--
 targets.go   |  4 ++--
 4 files changed, 63 insertions(+), 34 deletions(-)
 create mode 100644 main_test.go

diff --git a/main.go b/main.go
index c3b335c..dc85f53 100644
--- a/main.go
+++ b/main.go
@@ -30,14 +30,6 @@ func main() {
 	if *rate == 0 {
 		log.Fatal("rate can't be zero")
 	}
-	// Magic formula that assumes each client can
-	// sustain 200 RPS under normal circumstances
-	clients := make([]*Client, int(math.Ceil(float64(*rate)/200.0)))
-	ratePerClient := *rate / uint(len(clients))
-	for i := 0; i < len(clients); i++ {
-		clients[i] = NewClient(ratePerClient)
-	}
-
 	// Parse targets file
 	targets, err := NewTargetsFromFile(*targetsf)
 	if err != nil {
@@ -45,10 +37,8 @@ func main() {
 	}
 
 	// Parse ordering argument
-	random := false
 	if *ordering == "random" {
 		rand.Seed(time.Now().UnixNano())
-		random = true
 	} else if *ordering != "sequential" {
 		log.Fatalf("Unknown ordering %s", *ordering)
 	}
@@ -58,34 +48,46 @@ func main() {
 		log.Fatal("Duration provided is invalid")
 	}
 
-	hits := make(chan *http.Request, *rate*uint((*duration).Seconds()))
-	for i, idxs := 0, targets.Iter(random); i < cap(hits); i++ {
-		hits <- targets[idxs[i%len(idxs)]]
-	}
-	// Attack!
-	responses := make(chan *Response, cap(hits))
-	for _, client := range clients {
-		go client.Drill(hits, responses)
-	}
-	log.Printf("Vegeta is attacking ")
-	log.Printf("%d targets in %s order for %s with %d clients.\n", len(targets), *ordering, duration, len(clients))
-
+	// Parse reporter
 	var rep Reporter
 	switch *reporter {
 	case "text":
-		rep = NewTextReporter(len(responses))
+		rep = NewTextReporter()
 	default:
 		log.Println("reporter provided is not supported. using text")
-		rep = NewTextReporter(len(responses))
+		rep = NewTextReporter()
+	}
+
+	log.Printf("Vegeta is attacking %d targets in %s order for %s\n", len(targets), *ordering, *duration)
+	attack(targets, *ordering, *rate, *duration, rep)
+
+	// Report results!
+	if rep.Report(os.Stdout) != nil {
+		log.Fatal("Failed to report!")
+	}
+}
+
+func attack(targets Targets, ordering string, rate uint, duration time.Duration, rep Reporter) {
+	// Magic formula that assumes each client can
+	// sustain 200 RPS under normal circumstances
+	clients := make([]*Client, int(math.Ceil(float64(rate)/200.0)))
+	ratePerClient := rate / uint(len(clients))
+	for i := 0; i < len(clients); i++ {
+		clients[i] = NewClient(ratePerClient)
+	}
+
+	hits := make(chan *http.Request, rate*uint((duration).Seconds()))
+	defer close(hits)
+	for i, idxs := 0, targets.Iter(ordering); i < cap(hits); i++ {
+		hits <- targets[idxs[i%len(idxs)]]
+	}
+	responses := make(chan *Response, cap(hits))
+	defer close(responses)
+	for _, client := range clients {
+		go client.Drill(hits, responses) // Attack!
 	}
 	// Wait for all requests to finish
 	for i := 0; i < cap(responses); i++ {
 		rep.Add(<-responses)
 	}
-	close(hits)
-	close(responses)
-
-	if rep.Report(os.Stdout) != nil {
-		log.Fatal("Failed to report!")
-	}
 }
diff --git a/main_test.go b/main_test.go
new file mode 100644
index 0000000..982be0d
--- /dev/null
+++ b/main_test.go
@@ -0,0 +1,27 @@
+package main
+
+import (
+	"bytes"
+	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
+	"testing"
+	"time"
+)
+
+func TestAttackRate(t *testing.T) {
+	hitCount := uint64(0)
+	server := httptest.NewServer(
+		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			atomic.AddUint64(&hitCount, 1)
+		}),
+	)
+	targets, err := NewTargets(bytes.NewBufferString("GET " + server.URL + "\n"))
+	if err != nil {
+		t.Fatal(err)
+	}
+	attack(targets, "random", 50, 1*time.Second, NewTextReporter())
+	if hits := atomic.LoadUint64(&hitCount); hits != 50 {
+		t.Fatalf("Wrong number of hits: want %d, got %d\n", 50, hits)
+	}
+}
diff --git a/reporters.go b/reporters.go
index aee59de..a5ea100 100644
--- a/reporters.go
+++ b/reporters.go
@@ -17,8 +17,8 @@ type TextReporter struct {
 }
 
 // NewTextReporter initializes a TextReporter with n responses
-func NewTextReporter(n int) *TextReporter {
-	return &TextReporter{responses: make([]*Response, n)}
+func NewTextReporter() *TextReporter {
+	return &TextReporter{responses: make([]*Response, 0)}
 }
 
 // Add adds a response to be used in the report
diff --git a/targets.go b/targets.go
index 6fc3e7e..07d219e 100644
--- a/targets.go
+++ b/targets.go
@@ -49,8 +49,8 @@ func NewTargets(source io.Reader) (Targets, error) {
 	return targets, nil
 }
 
-func (t Targets) Iter(random bool) []int {
-	if random {
+func (t Targets) Iter(ordering string) []int {
+	if ordering == "random" {
 		return rand.Perm(len(t))
 	}
 	iter := make([]int, len(t))