Split Attack and Report into sub commands

attack.go

@@ -0,0 +1,67 @@
+package main
+
+import (
+	"fmt"
+	vegeta "github.com/tsenart/vegeta/lib"
+	"io"
+	"log"
+	"os"
+	"time"
+)
+
+const (
+	errRatePrefix        = "Rate: "
+	errDurationPrefix    = "Duration: "
+	errOutputFilePrefix  = "Output file: "
+	errTargetsFilePrefix = "Targets file: "
+	errOrderingPrefix    = "Ordering: "
+	errReportingPrefix   = "Reporting: "
+)
+
+// attack is the command function that validates the attack arguments, sets up the
+// required resources, launches the attack and reports 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")
+	}
+
+	if duration == 0 {
+		return fmt.Errorf(errDurationPrefix + "can't be zero")
+	}
+
+	targets, err := vegeta.NewTargetsFromFile(targetsf)
+	if err != nil {
+		return fmt.Errorf(errTargetsFilePrefix+"(%s): %s", targetsf, err)
+	}
+
+	switch ordering {
+	case "random":
+		targets.Shuffle(time.Now().UnixNano())
+	case "sequential":
+		break
+	default:
+		return fmt.Errorf(errOrderingPrefix+"`%s` is invalid", ordering)
+	}
+
+	var out io.Writer
+	switch output {
+	case "stdout":
+		out = os.Stdout
+	default:
+		file, err := os.Create(output)
+		if err != nil {
+			return fmt.Errorf(errOutputFilePrefix+"(%s): %s", output, err)
+		}
+		defer file.Close()
+		out = file
+	}
+
+	log.Printf("Vegeta is attacking %d targets in %s order for %s...\n", len(targets), ordering, duration)
+	results := vegeta.Attack(targets, rate, duration)
+	log.Println("Done!")
+	log.Printf("Writing results to '%s'...", output)
+	if err := results.WriteTo(out); err != nil {
+		return err
+	}
+	return nil
+}

attack_test.go

@@ -14,31 +14,31 @@ func init() {
 }
 
 func TestRateValidation(t *testing.T) {
-	rate, duration, targetsf, ordering, reporter, output := defaultArguments()
+	rate, duration, targetsf, ordering, output := defaultArguments()
 	rate = 0
 
-	err := run(rate, duration, targetsf, ordering, reporter, output)
+	err := attack(rate, duration, targetsf, ordering, output)
 	if err == nil || (err != nil && !strings.HasPrefix(err.Error(), errRatePrefix)) {
 		t.Errorf("Rate 0 shouldn't be valid: %s", err)
 	}
 }
 
 func TestDurationValidation(t *testing.T) {
-	rate, duration, targetsf, ordering, reporter, output := defaultArguments()
+	rate, duration, targetsf, ordering, output := defaultArguments()
 	duration = 0
 
-	err := run(rate, duration, targetsf, ordering, reporter, output)
+	err := attack(rate, duration, targetsf, ordering, output)
 	if err == nil || (err != nil && !strings.HasPrefix(err.Error(), errDurationPrefix)) {
 		t.Errorf("Duration 0 shouldn't be valid: %s", err)
 	}
 }
 
 func TestOutputValidation(t *testing.T) {
-	rate, duration, targetsf, ordering, reporter, _ := defaultArguments()
+	rate, duration, targetsf, ordering, _ := defaultArguments()
 
 	// Good cases
 	for _, output := range []string{"stdout", "/dev/null"} {
-		err := run(rate, duration, targetsf, ordering, reporter, output)
+		err := attack(rate, duration, targetsf, ordering, output)
 		if err != nil {
 			t.Errorf("Output file `%s` should be valid: %s", output, err)
 		}
@@ -46,46 +46,35 @@ func TestOutputValidation(t *testing.T) {
 
 	// Bad case
 	badOutput := ""
-	err := run(rate, duration, targetsf, ordering, reporter, badOutput)
+	err := attack(rate, duration, targetsf, ordering, badOutput)
 	if err == nil || (err != nil && !strings.HasPrefix(err.Error(), errOutputFilePrefix)) {
 		t.Errorf("Output file `%s` shouldn't be valid: %s", badOutput, err)
 	}
 }
 
-func TestReporter(t *testing.T) {
-	rate, duration, targetsf, ordering, _, output := defaultArguments()
-
-	for _, reporter := range []string{"text", "plot:timings"} {
-		err := run(rate, duration, targetsf, ordering, reporter, output)
-		if err != nil {
-			t.Errorf("Reporter `%s` shouldn't return an error: %s", reporter, err)
-		}
-	}
-}
-
 func TestTargetsValidation(t *testing.T) {
-	rate, duration, goodFile, ordering, reporter, output := defaultArguments()
+	rate, duration, goodFile, ordering, output := defaultArguments()
 
 	// Good case
-	err := run(rate, duration, goodFile, ordering, reporter, output)
+	err := attack(rate, duration, goodFile, ordering, output)
 	if err != nil {
 		t.Errorf("Targets file `%s` should be valid: %s", goodFile, err)
 	}
 
 	// Bad case
 	badFile := "randomInexistingFile12345.txt"
-	err = run(rate, duration, badFile, ordering, reporter, output)
+	err = attack(rate, duration, badFile, ordering, output)
 	if err == nil || (err != nil && !strings.HasPrefix(err.Error(), errTargetsFilePrefix)) {
 		t.Errorf("Targets file `%s` shouldn't be valid: %s", badFile, err)
 	}
 }
 
 func TestOrderingValidation(t *testing.T) {
-	rate, duration, targetsf, _, reporter, output := defaultArguments()
+	rate, duration, targetsf, _, output := defaultArguments()
 
 	// Good cases
 	for _, ordering := range []string{"random", "sequential"} {
-		err := run(rate, duration, targetsf, ordering, reporter, output)
+		err := attack(rate, duration, targetsf, ordering, output)
 		if err != nil {
 			t.Errorf("Ordering `%s` should be valid: %s", ordering, err)
 		}
@@ -93,12 +82,12 @@ func TestOrderingValidation(t *testing.T) {
 
 	// Bad case
 	badOrdering := "lolcat"
-	err := run(rate, duration, targetsf, badOrdering, reporter, output)
+	err := attack(rate, duration, targetsf, badOrdering, output)
 	if err == nil || (err != nil && !strings.HasPrefix(err.Error(), errOrderingPrefix)) {
 		t.Errorf("Ordering `%s` shouldn't be valid: %s", badOrdering, err)
 	}
 }
 
-func defaultArguments() (uint64, time.Duration, string, string, string, string) {
+func defaultArguments() (uint64, time.Duration, string, string, string) {
-	return uint64(1000), 5 * time.Millisecond, ".targets.txt", "random", "text", "/dev/null"
+	return uint64(1000), 5 * time.Millisecond, ".targets.txt", "random", "/dev/null"
 }

main.go

@@ -2,105 +2,49 @@ package main
 
 import (
 	"flag"
-	"fmt"
-	vegeta "github.com/tsenart/vegeta/lib"
-	"io"
 	"log"
-	"os"
 	"runtime"
 	"time"
 )
 
 func main() {
-	var (
+	// Global flags
-		rate     = flag.Uint64("rate", 50, "Requests per second")
+	cpus := flag.Int("cpus", runtime.NumCPU(), "Number of CPUs to use")
-		targetsf = flag.String("targets", "targets.txt", "Targets file")
-		ordering = flag.String("ordering", "random", "Attack ordering [sequential, random]")
-		duration = flag.Duration("duration", 10*time.Second, "Duration of the test")
-		reporter = flag.String("reporter", "text", "Reporter to use [text, json, plot:timings]")
-		output   = flag.String("output", "stdout", "Reporter output file")
-		cpus     = flag.Int("cpus", runtime.NumCPU(), "Number of CPUs to use")
-	)
 	flag.Parse()
-
-	if flag.NFlag() == 0 {
-		flag.Usage()
-		return
-	}
-
 	runtime.GOMAXPROCS(*cpus)
 
-	if err := run(*rate, *duration, *targetsf, *ordering, *reporter, *output); err != nil {
+	args := flag.Args()
-		log.Fatal(err)
+	if len(args) < 1 {
+		log.Fatal("Unspecified command")
 	}
-}
+	cmd, cmdf := args[0], flag.NewFlagSet(args[0], flag.ExitOnError)
 
-const (
+	switch cmd {
-	errRatePrefix        = "Rate: "
+	case "attack":
-	errDurationPrefix    = "Duration: "
+		rate := cmdf.Uint64("rate", 50, "Requests per second")
-	errOutputFilePrefix  = "Output file: "
+		targetsf := cmdf.String("targets", "targets.txt", "Targets file")
-	errTargetsFilePrefix = "Targets file: "
+		ordering := cmdf.String("ordering", "random", "Attack ordering [sequential, random]")
-	errOrderingPrefix    = "Ordering: "
+		duration := cmdf.Duration("duration", 10*time.Second, "Duration of the test")
-	errReportingPrefix   = "Reporting: "
+		output := cmdf.String("output", "stdout", "Vegeta Results file")
-)
 
-// run is an utility function that validates the attack arguments, sets up the
+		if err := cmdf.Parse(args[1:]); err != nil {
-// required resources, launches the attack and reports the results
+			log.Fatal(err)
-func run(rate uint64, duration time.Duration, targetsf, ordering, reporter, output string) error {
-	if rate == 0 {
-		return fmt.Errorf(errRatePrefix + "can't be zero")
-	}
-
-	if duration == 0 {
-		return fmt.Errorf(errDurationPrefix + "can't be zero")
-	}
-
-	var out io.Writer
-	switch output {
-	case "stdout":
-		out = os.Stdout
-	default:
-		file, err := os.Create(output)
-		if err != nil {
-			return fmt.Errorf(errOutputFilePrefix+"(%s): %s", output, err)
 		}
-		defer file.Close()
+		if err := attack(*rate, *duration, *targetsf, *ordering, *output); err != nil {
-		out = file
+			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")
 
-	var rep vegeta.Reporter
+		if err := cmdf.Parse(args[1:]); err != nil {
-	switch reporter {
+			log.Fatal(err)
-	case "text":
+		}
-		rep = vegeta.ReportText
+		if err := report(*reporter, *input, *output); err != nil {
-	case "json":
+			log.Fatal(err)
-		rep = vegeta.ReportJSON
+		}
-	case "plot:timings":
-		rep = vegeta.ReportTimingsPlot
 	default:
-		log.Println("Reporter provided is not supported. Using text")
+		log.Fatalf("Unknown command: %s", cmd)
-		rep = vegeta.ReportText
 	}
-
-	targets, err := vegeta.NewTargetsFromFile(targetsf)
-	if err != nil {
-		return fmt.Errorf(errTargetsFilePrefix+"(%s): %s", targetsf, err)
-	}
-
-	switch ordering {
-	case "random":
-		targets.Shuffle(time.Now().UnixNano())
-	case "sequential":
-		break
-	default:
-		return fmt.Errorf(errOrderingPrefix+"`%s` is invalid", ordering)
-	}
-
-	log.Printf("Vegeta is attacking %d targets in %s order for %s...\n", len(targets), ordering, duration)
-	results := vegeta.Attack(targets, rate, duration)
-	log.Println("Done!")
-	log.Printf("Writing report to '%s'...", output)
-	if err = rep(results, out); err != nil {
-		return fmt.Errorf(errReportingPrefix+"%s", err)
-	}
-	return nil
 }

report.go

@@ -0,0 +1,59 @@
+package main
+
+import (
+	vegeta "github.com/tsenart/vegeta/lib"
+	"io"
+	"log"
+	"os"
+)
+
+func report(reporter, input, output string) error {
+	var rep vegeta.Reporter
+	switch reporter {
+	case "text":
+		rep = vegeta.ReportText
+	case "json":
+		rep = vegeta.ReportJSON
+	case "plot:timings":
+		rep = vegeta.ReportTimingsPlot
+	default:
+		log.Println("Reporter provided is not supported. Using text")
+		rep = vegeta.ReportText
+	}
+
+	var in io.Reader
+	switch input {
+	case "stdin":
+		in = os.Stdin
+	default:
+		file, err := os.Open(input)
+		if err != nil {
+			return err
+		}
+		defer file.Close()
+		in = file
+	}
+
+	var out io.Writer
+	switch output {
+	case "stdout":
+		out = os.Stdout
+	default:
+		file, err := os.Create(output)
+		if err != nil {
+			return err
+		}
+		defer file.Close()
+		out = file
+	}
+
+	results := vegeta.Results{}
+	if err := results.ReadFrom(in); err != nil {
+		return err
+	}
+	if err := rep(results, out); err != nil {
+		return err
+	}
+
+	return nil
+}