From 7cd13327c62f78cb475af50d4763386108ae29dd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Tom=C3=A1s=20Senart?= Date: Wed, 11 Sep 2013 14:10:45 +0100 Subject: [PATCH] Default targets to stdin and abstract file handling --- attack.go | 26 +++++++++++--------------- attack_test.go | 19 ------------------- file.go | 26 ++++++++++++++++++++++++++ lib/targets.go | 15 ++------------- lib/targets_test.go | 4 ++-- report.go | 33 +++++++++------------------------ 6 files changed, 50 insertions(+), 73 deletions(-) create mode 100644 file.go diff --git a/attack.go b/attack.go index 5752bee..8d7e83b 100644 --- a/attack.go +++ b/attack.go @@ -4,16 +4,14 @@ import ( "flag" "fmt" vegeta "github.com/tsenart/vegeta/lib" - "io" "log" - "os" "time" ) 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") + targetsf := fs.String("targets", "stdin", "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", "Output file") @@ -35,7 +33,12 @@ func attack(rate uint64, duration time.Duration, targetsf, ordering, output stri return fmt.Errorf(errDurationPrefix + "can't be zero") } - targets, err := vegeta.NewTargetsFromFile(targetsf) + in, err := file(targetsf, false) + if err != nil { + return fmt.Errorf(errTargetsFilePrefix+"(%s): %s", targetsf, err) + } + defer in.Close() + targets, err := vegeta.NewTargetsFrom(in) if err != nil { return fmt.Errorf(errTargetsFilePrefix+"(%s): %s", targetsf, err) } @@ -49,18 +52,11 @@ func attack(rate uint64, duration time.Duration, targetsf, ordering, output stri 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 + out, err := file(output, true) + if err != nil { + return fmt.Errorf(errOutputFilePrefix+"(%s): %s", output, err) } + defer out.Close() log.Printf("Vegeta is attacking %d targets in %s order for %s...\n", len(targets), ordering, duration) results := vegeta.Attack(targets, rate, duration) diff --git a/attack_test.go b/attack_test.go index 824386d..f9b1216 100644 --- a/attack_test.go +++ b/attack_test.go @@ -33,25 +33,6 @@ func TestDurationValidation(t *testing.T) { } } -func TestOutputValidation(t *testing.T) { - rate, duration, targetsf, ordering, _ := defaultArguments() - - // Good cases - for _, output := range []string{"stdout", "/dev/null"} { - err := attack(rate, duration, targetsf, ordering, output) - if err != nil { - t.Errorf("Output file `%s` should be valid: %s", output, err) - } - } - - // Bad case - 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 TestTargetsValidation(t *testing.T) { rate, duration, goodFile, ordering, output := defaultArguments() diff --git a/file.go b/file.go new file mode 100644 index 0000000..539c313 --- /dev/null +++ b/file.go @@ -0,0 +1,26 @@ +package main + +import ( + "os" +) + +func file(filename string, create bool) (*os.File, error) { + switch filename { + case "stdin": + return os.Stdin, nil + case "stdout": + return os.Stdout, nil + default: + var file *os.File + var err error + if create { + file, err = os.Create(filename) + } else { + file, err = os.Open(filename) + } + if err != nil { + return nil, err + } + return file, nil + } +} diff --git a/lib/targets.go b/lib/targets.go index 467b899..9864f86 100644 --- a/lib/targets.go +++ b/lib/targets.go @@ -6,25 +6,14 @@ import ( "io" "math/rand" "net/http" - "os" "strings" ) // Targets represents the http.Requests which will be issued during the test type Targets []*http.Request -// NewTargetsFromFile reads and parses targets from a text file -func NewTargetsFromFile(filename string) (Targets, error) { - file, err := os.Open(filename) - if err != nil { - return Targets{}, err - } - defer file.Close() - return readTargets(file) -} - -// readTargets reads targets out of a line separated source skipping empty lines -func readTargets(source io.Reader) (Targets, error) { +// NewTargetsFrom reads targets out of a line separated source skipping empty lines +func NewTargetsFrom(source io.Reader) (Targets, error) { scanner := bufio.NewScanner(source) lines := make([]string, 0) for scanner.Scan() { diff --git a/lib/targets_test.go b/lib/targets_test.go index 6b7e99e..9337928 100644 --- a/lib/targets_test.go +++ b/lib/targets_test.go @@ -6,9 +6,9 @@ import ( "testing" ) -func TestReadTargets(t *testing.T) { +func TestNewTargetsFrom(t *testing.T) { lines := bytes.NewBufferString("GET http://lolcathost:9999/\n\n // HEAD http://lolcathost.com this is a comment \nHEAD http://lolcathost:9999/\n") - targets, err := readTargets(lines) + targets, err := NewTargetsFrom(lines) if err != nil { t.Fatalf("Couldn't parse valid source: %s", err) } diff --git a/report.go b/report.go index d0a5499..e428e7b 100644 --- a/report.go +++ b/report.go @@ -3,9 +3,7 @@ package main import ( "flag" vegeta "github.com/tsenart/vegeta/lib" - "io" "log" - "os" ) func reportCmd(args []string) command { @@ -36,36 +34,23 @@ func report(reporter, input, output string) error { 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 + in, err := file(input, false) + if err != nil { + return err } + defer in.Close() - 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 + out, err := file(output, true) + if err != nil { + return err } + defer out.Close() results := vegeta.Results{} if err := results.ReadFrom(in); err != nil { return err } + if err := rep(results, out); err != nil { return err }