This commit adds some tests

Author: codesenberg

.gitignore

@@ -22,3 +22,4 @@ _testmain.go
 *.exe
 *.test
 *.prof
+cover.*

args_parser.go

@@ -0,0 +1,86 @@
+package main
+
+import (
+	"errors"
+	"flag"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"os"
+	"time"
+)
+
+var (
+	errNoURL       = errors.New("No URL supplied")
+	errTooManyArgs = errors.New("Too many arguments are supplied")
+
+	emptyConf = config{}
+	parser    = newDefaultParser()
+)
+
+type defaultParser struct {
+	fs *flag.FlagSet
+
+	numReqs   *nullableUint64
+	duration  *nullableDuration
+	headers   *headersList
+	numConns  uint64
+	timeout   time.Duration
+	latencies bool
+	method    string
+	body      string
+}
+
+func newDefaultParser() *defaultParser {
+	dp := new(defaultParser)
+	dp.fs = flag.NewFlagSet(programName, flag.ContinueOnError)
+	dp.numReqs = new(nullableUint64)
+	dp.duration = new(nullableDuration)
+	dp.headers = new(headersList)
+	dp.fs.Uint64Var(&dp.numConns, "c", defaultNumberOfConns, "Maximum number of concurrent connections")
+	dp.fs.DurationVar(&dp.timeout, "timeout", defaultTimeout, "Socket/request timeout")
+	dp.fs.BoolVar(&dp.latencies, "latencies", false, "Print latency statistics")
+	dp.fs.StringVar(&dp.method, "m", "GET", "Request method")
+	dp.fs.StringVar(&dp.body, "data", "", "Request body")
+	dp.fs.Var(dp.headers, "H", "HTTP headers to use")
+	dp.fs.Var(dp.numReqs, "n", "Number of requests")
+	dp.fs.Var(dp.duration, "d", "Duration of test")
+	return dp
+}
+
+func (p *defaultParser) usage(out io.Writer) {
+	fmt.Fprintln(out, programName, "<options> <url>")
+	p.fs.SetOutput(out)
+	p.fs.PrintDefaults()
+}
+
+func (p *defaultParser) parse(args []string) (config, error) {
+	err := p.parseArgs(args)
+	if err != nil {
+		return emptyConf, err
+	}
+	if p.fs.NArg() == 0 {
+		return emptyConf, errNoURL
+	}
+	if p.fs.NArg() > 1 {
+		return emptyConf, errTooManyArgs
+	}
+	return config{
+		numConns:       p.numConns,
+		numReqs:        p.numReqs.val,
+		duration:       p.duration.val,
+		url:            p.fs.Arg(0),
+		headers:        p.headers,
+		timeout:        p.timeout,
+		method:         p.method,
+		body:           p.body,
+		printLatencies: p.latencies,
+	}, nil
+}
+
+func (p *defaultParser) parseArgs(args []string) error {
+	p.fs.SetOutput(ioutil.Discard)
+	err := p.fs.Parse(args[1:])
+	p.fs.SetOutput(os.Stdout)
+	return err
+}

args_parser_test.go

@@ -0,0 +1,54 @@
+package main
+
+import (
+	"bytes"
+	"testing"
+)
+
+func TestArgsParsing(t *testing.T) {
+	expectations := []struct {
+		p   *defaultParser
+		in  []string
+		out error
+	}{
+		{
+			newDefaultParser(),
+			[]string{programName},
+			errNoURL,
+		},
+		{
+			newDefaultParser(),
+			[]string{programName, "http://google.com", "http://yahoo.com"},
+			errTooManyArgs,
+		},
+		{
+			newDefaultParser(),
+			[]string{programName, "http://google.com"},
+			nil,
+		},
+	}
+	for _, e := range expectations {
+		if _, err := e.p.parse(e.in); err != e.out {
+			t.Log(err, e.out)
+			t.Fail()
+		}
+	}
+}
+
+func TestUnspecifiedArgParsing(t *testing.T) {
+	p := newDefaultParser()
+	args := []string{programName, "--someunspecifiedflag"}
+	_, err := p.parse(args)
+	if err == nil {
+		t.Fail()
+	}
+}
+
+func TestUsagePrinting(t *testing.T) {
+	b := new(bytes.Buffer)
+	p := newDefaultParser()
+	p.usage(b)
+	if b.Len() == 0 {
+		t.Fail()
+	}
+}

bombardier.go

@@ -1,8 +1,8 @@
 package main
 
 import (
-	"flag"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"os"
 	"sync"
@@ -13,20 +13,16 @@ import (
 	"github.com/valyala/fasthttp"
 )
 
-const (
-	maxRps           = 10000000
-	requestsInterval = 100 * time.Millisecond
-)
-
 type bombardier struct {
 	conf           config
 	requestHeaders *fasthttp.RequestHeader
 	barrier        completionBarrier
 
-	bytesWritten int64
-	timeTaken    time.Duration
-	latencies    *stats
-	requests     *stats
+	bytesTotal int64
+	bytesData  int64
+	timeTaken  time.Duration
+	latencies  *stats
+	requests   *stats
 
 	client *fasthttp.Client
 	done   chan bool
@@ -46,6 +42,9 @@ type bombardier struct {
 
 	// Progress bar
 	bar *pb.ProgressBar
+
+	// Output
+	out io.Writer
 }
 
 func newBombardier(c config) (*bombardier, error) {
@@ -65,10 +64,11 @@ func newBombardier(c config) (*bombardier, error) {
 		b.bar = pb.New(int(b.conf.duration.Seconds()))
 		b.bar.ShowCounters = false
 		b.bar.ShowPercent = false
-		b.barrier = newTimedCompletionBarrier(int(c.numConns), *c.duration, func() {
+		b.barrier = newTimedCompletionBarrier(c.numConns, tickDuration, *c.duration, func() {
 			b.bar.Increment()
 		})
 	}
+	b.out = os.Stdout
 	b.client = &fasthttp.Client{
 		MaxConnsPerHost: int(c.numConns),
 	}
@@ -89,15 +89,16 @@ func (b *bombardier) prepareRequest() (*fasthttp.Request, *fasthttp.Response) {
 	return req, resp
 }
 
-func (b *bombardier) fireRequest(req *fasthttp.Request, resp *fasthttp.Response) (bytesWritten int64, code int, msTaken uint64) {
+func (b *bombardier) fireRequest(req *fasthttp.Request, resp *fasthttp.Response) (bytesData, bytesTotal int64, code int, msTaken uint64) {
 	start := time.Now()
 	err := b.client.DoTimeout(req, resp, b.conf.timeout)
 	if err != nil {
 		code = 0
 	} else {
 		code = resp.StatusCode()
 	}
-	bytesWritten, _ = resp.WriteTo(ioutil.Discard)
+	bytesData = int64(len(resp.Body()))
+	bytesTotal, _ = resp.WriteTo(ioutil.Discard)
 	msTaken = uint64(time.Since(start).Nanoseconds() / 1000)
 	return
 }
@@ -107,9 +108,10 @@ func (b *bombardier) releaseRequest(req *fasthttp.Request, resp *fasthttp.Respon
 	fasthttp.ReleaseResponse(resp)
 }
 
-func (b *bombardier) writeStatistics(bytesWritten int64, code int, msTaken uint64) {
+func (b *bombardier) writeStatistics(bytesData, bytesTotal int64, code int, msTaken uint64) {
 	b.latencies.record(msTaken)
-	atomic.AddInt64(&b.bytesWritten, bytesWritten)
+	atomic.AddInt64(&b.bytesTotal, bytesTotal)
+	atomic.AddInt64(&b.bytesData, bytesData)
 	b.rpl.Lock()
 	b.reqs++
 	b.rpl.Unlock()
@@ -133,12 +135,12 @@ func (b *bombardier) writeStatistics(bytesWritten int64, code int, msTaken uint6
 	atomic.AddUint64(counter, 1)
 }
 
-func (b *bombardier) Worker() {
+func (b *bombardier) worker() {
 	for b.barrier.grabWork() {
 		req, resp := b.prepareRequest()
-		bytesWritten, code, msTaken := b.fireRequest(req, resp)
+		bytesData, bytesTotal, code, msTaken := b.fireRequest(req, resp)
 		b.releaseRequest(req, resp)
-		b.writeStatistics(bytesWritten, code, msTaken)
+		b.writeStatistics(bytesData, bytesTotal, code, msTaken)
 		b.barrier.jobDone()
 	}
 }
@@ -174,94 +176,88 @@ func (b *bombardier) bombard() {
 	bombardmentBegin := time.Now()
 	b.start = time.Now()
 	for i := uint64(0); i < b.conf.numConns; i++ {
-		go b.Worker()
+		go b.worker()
 	}
 	go b.rateMeter()
 	b.barrier.wait()
 	b.timeTaken = time.Since(bombardmentBegin)
 	b.done <- true
 	<-b.done
-	b.bar.FinishPrint("Done!")
+	b.bar.Finish()
+	fmt.Fprintln(b.out, "Done!")
 }
 
 func (b *bombardier) throughput() float64 {
-	return float64(b.bytesWritten) / b.timeTaken.Seconds()
+	return float64(b.bytesTotal) / b.timeTaken.Seconds()
 }
 
 func (b *bombardier) printIntro() {
 	if b.conf.testType == counted {
-		fmt.Printf("Bombarding %v with %v requests using %v connections\n",
+		fmt.Fprintf(b.out, "Bombarding %v with %v requests using %v connections\n",
 			b.conf.url, *b.conf.numReqs, b.conf.numConns)
 	} else if b.conf.testType == timed {
-		fmt.Printf("Bombarding %v for %v using %v connections\n",
+		fmt.Fprintf(b.out, "Bombarding %v for %v using %v connections\n",
 			b.conf.url, *b.conf.duration, b.conf.numConns)
 	}
 }
 
 func (b *bombardier) printLatencyStats() {
 	percentiles := []float64{50.0, 75.0, 90.0, 99.0}
-	fmt.Println("  Latency Distribution")
+	fmt.Fprintln(b.out, "  Latency Distribution")
 	for i := 0; i < len(percentiles); i++ {
 		p := percentiles[i]
 		n := b.latencies.percentile(p)
-		fmt.Printf("     %2.0f%% %10s", p, formatUnits(float64(n), timeUnitsUs, 2))
-		fmt.Printf("\n")
+		fmt.Fprintf(b.out, "     %2.0f%% %10s", p, formatUnits(float64(n), timeUnitsUs, 2))
+		fmt.Fprintf(b.out, "\n")
 	}
 }
 
 func (b *bombardier) printStats() {
-	fmt.Printf("%10v %10v %10v %10v\n", "Statistics", "Avg", "Stdev", "Max")
-	fmt.Println(rpsString(b.requests))
-	fmt.Println(latenciesString(b.latencies))
-	if *latencies {
+	fmt.Fprintf(b.out, "%10v %10v %10v %10v\n", "Statistics", "Avg", "Stdev", "Max")
+	fmt.Fprintln(b.out, rpsString(b.requests))
+	fmt.Fprintln(b.out, latenciesString(b.latencies))
+	if b.conf.printLatencies {
 		b.printLatencyStats()
 	}
-	fmt.Println("  HTTP codes:")
-	fmt.Printf("    1xx - %v, 2xx - %v, 3xx - %v, 4xx - %v, 5xx - %v\n",
+	fmt.Fprintln(b.out, "  HTTP codes:")
+	fmt.Fprintf(b.out, "    1xx - %v, 2xx - %v, 3xx - %v, 4xx - %v, 5xx - %v\n",
 		b.req1xx, b.req2xx, b.req3xx, b.req4xx, b.req5xx)
-	fmt.Printf("    errored - %v\n", b.errored)
-	fmt.Printf("  %-10v %10v/s\n", "Throughput:", formatBinary(b.throughput()))
+	fmt.Fprintf(b.out, "    errored - %v\n", b.errored)
+	fmt.Fprintf(b.out, "  %-10v %10v/s\n", "Throughput:", formatBinary(b.throughput()))
+}
+
+func (b *bombardier) redirectOutputTo(out io.Writer) {
+	b.bar.Output = out
+	b.out = out
+}
+
+func (b *bombardier) disableOutput() {
+	b.redirectOutputTo(ioutil.Discard)
+	b.bar.NotPrint = true
 }
 
-var (
-	numReqs   = new(nullableUint64)
-	duration  = new(nullableDuration)
-	headers   = new(headersList)
-	numConns  = flag.Uint64("c", 200, "Maximum number of concurrent connections")
-	timeout   = flag.Duration("timeout", 2*time.Second, "Socket/request timeout")
-	latencies = flag.Bool("latencies", false, "Print latency statistics")
-	method    = flag.String("m", "GET", "Request method")
-	body      = flag.String("data", "", "Request body")
+const (
+	maxRps           = 10000000
+	requestsInterval = 100 * time.Millisecond
+	tickDuration     = 1 * time.Second
+	defaultTimeout   = 2 * time.Second
+
+	programName = "bombardier"
+
+	EXIT_FAILURE = 1
 )
 
 func main() {
-	flag.Var(headers, "H", "HTTP headers to use")
-	flag.Var(numReqs, "n", "Number of requests")
-	flag.Var(duration, "d", "Duration of test")
-	flag.Parse()
-	if flag.NArg() == 0 {
-		fmt.Println("No URL supplied")
-		flag.Usage()
-		os.Exit(1)
-	}
-	if flag.NArg() > 1 {
-		fmt.Println("Too many arguments are supplied")
-		os.Exit(1)
+	config, err := parser.parse(os.Args)
+	if err != nil {
+		fmt.Println(err)
+		parser.usage(os.Stdout)
+		os.Exit(EXIT_FAILURE)
 	}
-	bombardier, err := newBombardier(
-		config{
-			numConns: *numConns,
-			numReqs:  numReqs.val,
-			duration: duration.val,
-			url:      flag.Arg(0),
-			headers:  headers,
-			timeout:  *timeout,
-			method:   *method,
-			body:     *body,
-		})
+	bombardier, err := newBombardier(config)
 	if err != nil {
 		fmt.Println(err)
-		os.Exit(1)
+		os.Exit(EXIT_FAILURE)
 	}
 	bombardier.bombard()
 	bombardier.printStats()