senaduka/old-vegeta
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Split Attack and Report into sub commands

Browse Source
This commit is contained in:
Tomás Senart
2013-09-10 21:47:37 +01:00
parent 6348a1b7df
commit 4bbc7a0296
4 changed files with 170 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files

67
attack.go Normal file
View File

@@ -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
}

41
main_test.go → attack_test.go
View File

@@ -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) {
return uint64(1000), 5 * time.Millisecond, ".targets.txt", "random", "text", "/dev/null"
func defaultArguments() (uint64, time.Duration, string, string, string) {
return uint64(1000), 5 * time.Millisecond, ".targets.txt", "random", "/dev/null"
}

114
main.go
View File

@@ -2,105 +2,49 @@ package main
import (
"flag"
"fmt"
vegeta "github.com/tsenart/vegeta/lib"
"io"
"log"
"os"
"runtime"
"time"
)
func main() {
var (
rate = flag.Uint64("rate", 50, "Requests per second")
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")
)
// Global flags
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 {
log.Fatal(err)
args := flag.Args()
if len(args) < 1 {
log.Fatal("Unspecified command")
}
}
cmd, cmdf := args[0], flag.NewFlagSet(args[0], flag.ExitOnError)
const (
errRatePrefix = "Rate: "
errDurationPrefix = "Duration: "
errOutputFilePrefix = "Output file: "
errTargetsFilePrefix = "Targets file: "
errOrderingPrefix = "Ordering: "
errReportingPrefix = "Reporting: "
)
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")
// run is an utility function that validates the attack arguments, sets up the
// required resources, launches the attack and reports the results
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)
if err := cmdf.Parse(args[1:]); err != nil {
log.Fatal(err)
}
defer file.Close()
out = file
}
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")
var rep vegeta.Reporter
switch reporter {
case "text":
rep = vegeta.ReportText
case "json":
rep = vegeta.ReportJSON
case "plot:timings":
rep = vegeta.ReportTimingsPlot
if err := cmdf.Parse(args[1:]); err != nil {
log.Fatal(err)
}
if err := report(*reporter, *input, *output); err != nil {
log.Fatal(err)
}
default:
log.Println("Reporter provided is not supported. Using text")
rep = vegeta.ReportText
log.Fatalf("Unknown command: %s", cmd)
}
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
}

59
report.go Normal file
View File

@@ -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
}