Default targets to stdin and abstract file handling

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)

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()
 

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

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() {

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

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
 	}