From e61ded8e5b37d3c37efd4f17fd4b9f2b08c456b1 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Tom=C3=A1s=20Senart?= <tomas@soundcloud.com>
Date: Wed, 11 Sep 2013 00:10:35 +0100
Subject: [PATCH] Abstract command building and setup with closures

---
 attack.go | 35 +++++++++++++++++++++++++----------
 main.go   | 48 ++++++++++++++----------------------------------
 report.go | 15 +++++++++++++++
 3 files changed, 54 insertions(+), 44 deletions(-)

diff --git a/attack.go b/attack.go
index b10e5ab..bceabb7 100644
--- a/attack.go
+++ b/attack.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"flag"
 	"fmt"
 	vegeta "github.com/tsenart/vegeta/lib"
 	"io"
@@ -9,17 +10,22 @@ import (
 	"time"
 )
 
-const (
-	errRatePrefix        = "Rate: "
-	errDurationPrefix    = "Duration: "
-	errOutputFilePrefix  = "Output file: "
-	errTargetsFilePrefix = "Targets file: "
-	errOrderingPrefix    = "Ordering: "
-	errReportingPrefix   = "Reporting: "
-)
+func attackCmd(args []string) command {
+	fs := flag.NewFlagSet("attack", flag.ExitOnError)
+	rate := fs.Uint64("rate", 50, "Requests per second")
+	targetsf := fs.String("targets", "targets.txt", "Targets file")
+	ordering := fs.String("ordering", "random", "Attack ordering [sequential, random]")
+	duration := fs.Duration("duration", 10*time.Second, "Duration of the test")
+	output := fs.String("output", "stdout", "Vegeta Results file")
+	fs.Parse(args)
 
-// attack is the command function that validates the attack arguments, sets up the
-// required resources, launches the attack and reports the results
+	return func() error {
+		return attack(*rate, *duration, *targetsf, *ordering, *output)
+	}
+}
+
+// attack validates the attack arguments, sets up the
+// required resources, launches the attack and writes the results
 func attack(rate uint64, duration time.Duration, targetsf, ordering, output string) error {
 	if rate == 0 {
 		return fmt.Errorf(errRatePrefix + "can't be zero")
@@ -65,3 +71,12 @@ func attack(rate uint64, duration time.Duration, targetsf, ordering, output stri
 	}
 	return nil
 }
+
+const (
+	errRatePrefix        = "Rate: "
+	errDurationPrefix    = "Duration: "
+	errOutputFilePrefix  = "Output file: "
+	errTargetsFilePrefix = "Targets file: "
+	errOrderingPrefix    = "Ordering: "
+	errReportingPrefix   = "Reporting: "
+)
diff --git a/main.go b/main.go
index 31b62c9..bd73933 100644
--- a/main.go
+++ b/main.go
@@ -4,47 +4,27 @@ import (
 	"flag"
 	"log"
 	"runtime"
-	"time"
 )
 
+// command is a closure function which each command constructor
+// builds and returns
+type command func() error
+
 func main() {
+	commands := map[string]func([]string) command{
+		"attack": attackCmd,
+		"report": reportCmd,
+	}
 	// Global flags
 	cpus := flag.Int("cpus", runtime.NumCPU(), "Number of CPUs to use")
 	flag.Parse()
+	args := flag.Args()
+
 	runtime.GOMAXPROCS(*cpus)
 
-	args := flag.Args()
-	if len(args) < 1 {
-		log.Fatal("Unspecified command")
-	}
-	cmd, cmdf := args[0], flag.NewFlagSet(args[0], flag.ExitOnError)
-
-	switch cmd {
-	case "attack":
-		rate := cmdf.Uint64("rate", 50, "Requests per second")
-		targetsf := cmdf.String("targets", "targets.txt", "Targets file")
-		ordering := cmdf.String("ordering", "random", "Attack ordering [sequential, random]")
-		duration := cmdf.Duration("duration", 10*time.Second, "Duration of the test")
-		output := cmdf.String("output", "stdout", "Vegeta Results file")
-
-		if err := cmdf.Parse(args[1:]); err != nil {
-			log.Fatal(err)
-		}
-		if err := attack(*rate, *duration, *targetsf, *ordering, *output); err != nil {
-			log.Fatal(err)
-		}
-	case "report":
-		reporter := cmdf.String("reporter", "text", "Reporter [text, json, plot:timings]")
-		input := cmdf.String("input", "stdin", "Vegeta Results file")
-		output := cmdf.String("output", "stdout", "Output file")
-
-		if err := cmdf.Parse(args[1:]); err != nil {
-			log.Fatal(err)
-		}
-		if err := report(*reporter, *input, *output); err != nil {
-			log.Fatal(err)
-		}
-	default:
-		log.Fatalf("Unknown command: %s", cmd)
+	if cmd, ok := commands[args[0]]; !ok {
+		log.Fatalf("Unknown command: %s", args[0])
+	} else if err := cmd(args[1:])(); err != nil {
+		log.Fatal(err)
 	}
 }
diff --git a/report.go b/report.go
index f99eb6d..69bf5e4 100644
--- a/report.go
+++ b/report.go
@@ -1,12 +1,27 @@
 package main
 
 import (
+	"flag"
 	vegeta "github.com/tsenart/vegeta/lib"
 	"io"
 	"log"
 	"os"
 )
 
+func reportCmd(args []string) command {
+	fs := flag.NewFlagSet("report", flag.ExitOnError)
+	reporter := fs.String("reporter", "text", "Reporter [text, json, plot:timings]")
+	input := fs.String("input", "stdin", "Vegeta Results file")
+	output := fs.String("output", "stdout", "Output file")
+	fs.Parse(args)
+
+	return func() error {
+		return report(*reporter, *input, *output)
+	}
+}
+
+// report validates the report arguments, sets up the required resources
+// and writes the report
 func report(reporter, input, output string) error {
 	var rep vegeta.Reporter
 	switch reporter {