From 9f478d96fad42b3bca9cd8b1e15ea6e7ca5a824f Mon Sep 17 00:00:00 2001 From: nakabonne Date: Sat, 12 Sep 2020 21:13:56 +0900 Subject: [PATCH] Add fake attacker --- attacker/attacker.go | 19 ++++++++++++++----- attacker/attacker_test.go | 28 ++++++++++++++++++++++++++++ attacker/fake_attacker.go | 18 ++++++++++++++++++ go.mod | 1 + go.sum | 21 +++++++++------------ main.go | 1 + 6 files changed, 71 insertions(+), 17 deletions(-) create mode 100644 attacker/attacker_test.go create mode 100644 attacker/fake_attacker.go diff --git a/attacker/attacker.go b/attacker/attacker.go index e5512e7..0adcd50 100644 --- a/attacker/attacker.go +++ b/attacker/attacker.go @@ -16,6 +16,11 @@ const ( defaultMethod = http.MethodGet ) +type Attacker interface { + Attack(vegeta.Targeter, vegeta.Pacer, time.Duration, string) <-chan *vegeta.Result + Stop() +} + // Options provides optional settings to attack. type Options struct { Rate int @@ -23,8 +28,7 @@ type Options struct { Method string Body []byte Header http.Header - // TODO: Make attacker pluggable to make it easy unit tests. - //Attacker func(vegeta.Targeter, vegeta.Pacer, time.Duration, string) <-chan *vegeta.Result + Attacker Attacker } // Result contains the results of a single Target hit. @@ -36,6 +40,9 @@ type Result struct { // Results are sent to the given channel as soon as they arrive. // When the attack is over, it gives back final statistics. func Attack(ctx context.Context, target string, resCh chan *Result, opts Options) *Metrics { + if target == "" { + return nil + } if opts.Rate == 0 { opts.Rate = defaultRate } @@ -45,6 +52,9 @@ func Attack(ctx context.Context, target string, resCh chan *Result, opts Options if opts.Method == "" { opts.Method = defaultMethod } + if opts.Attacker == nil { + opts.Attacker = vegeta.NewAttacker() + } rate := vegeta.Rate{Freq: opts.Rate, Per: time.Second} targeter := vegeta.NewStaticTargeter(vegeta.Target{ @@ -53,14 +63,13 @@ func Attack(ctx context.Context, target string, resCh chan *Result, opts Options Body: opts.Body, Header: opts.Header, }) - attacker := vegeta.NewAttacker() var metrics vegeta.Metrics - for res := range attacker.Attack(targeter, rate, opts.Duration, "main") { + for res := range opts.Attacker.Attack(targeter, rate, opts.Duration, "main") { select { case <-ctx.Done(): - attacker.Stop() + opts.Attacker.Stop() default: resCh <- &Result{Latency: res.Latency} metrics.Add(res) diff --git a/attacker/attacker_test.go b/attacker/attacker_test.go new file mode 100644 index 0000000..94183d4 --- /dev/null +++ b/attacker/attacker_test.go @@ -0,0 +1,28 @@ +package attacker + +import ( + "context" + "testing" + + "github.com/stretchr/testify/assert" +) + +func TestAttack(t *testing.T) { + tests := []struct { + name string + target string + resCh chan *Result + opts Options + want *Metrics + }{ + // TODO: Add test cases. + } + ctx := context.Background() + for _, tt := range tests { + tt.opts.Attacker = &fakeAttacker{} + t.Run(tt.name, func(t *testing.T) { + got := Attack(ctx, tt.target, tt.resCh, tt.opts) + assert.Equal(t, got, tt.want) + }) + } +} diff --git a/attacker/fake_attacker.go b/attacker/fake_attacker.go new file mode 100644 index 0000000..3e67e8b --- /dev/null +++ b/attacker/fake_attacker.go @@ -0,0 +1,18 @@ +package attacker + +import ( + "time" + + vegeta "github.com/tsenart/vegeta/v12/lib" +) + +type fakeAttacker struct { +} + +func (f *fakeAttacker) Attack(vegeta.Targeter, vegeta.Pacer, time.Duration, string) <-chan *vegeta.Result { + return nil +} + +func (f *fakeAttacker) Stop() { + return +} diff --git a/go.mod b/go.mod index 7f7e481..25e6ca0 100644 --- a/go.mod +++ b/go.mod @@ -8,5 +8,6 @@ require ( github.com/mattn/go-colorable v0.1.7 // indirect github.com/mum4k/termdash v0.12.2 github.com/spf13/pflag v1.0.5 + github.com/stretchr/testify v1.6.1 github.com/tsenart/vegeta/v12 v12.8.3 ) diff --git a/go.sum b/go.sum index 288f42d..69fcb03 100644 --- a/go.sum +++ b/go.sum @@ -3,8 +3,8 @@ github.com/alecthomas/jsonschema v0.0.0-20180308105923-f2c93856175a/go.mod h1:qp github.com/bmizerany/perks v0.0.0-20141205001514-d9a9656a3a4b h1:AP/Y7sqYicnjGDfD5VcY4CIfh1hRXBUavxrvELjTiOE= github.com/bmizerany/perks v0.0.0-20141205001514-d9a9656a3a4b/go.mod h1:ac9efd0D1fsDb3EJvhqgXRbFx7bs2wqZ10HQPeU8U/Q= github.com/c2h5oh/datasize v0.0.0-20171227191756-4eba002a5eae/go.mod h1:S/7n9copUssQ56c7aAgHqftWO4LTf4xY6CGWt8Bc+3M= -github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= -github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8= +github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/dgryski/go-gk v0.0.0-20140819190930-201884a44051 h1:ByJUvQYyTtNNCVfYNM48q6uYUT4fAlN0wNmd3th4BSo= github.com/dgryski/go-gk v0.0.0-20140819190930-201884a44051/go.mod h1:qm+vckxRlDt0aOla0RYJJVeqHZlWfOm2UIxHaqPB46E= github.com/dgryski/go-lttb v0.0.0-20180810165845-318fcdf10a77/go.mod h1:Va5MyIzkU0rAM92tn3hb3Anb7oz7KcnixF49+2wOMe4= @@ -36,8 +36,6 @@ github.com/k0kubun/colorstring v0.0.0-20150214042306-9440f1994b88 h1:uC1QfSlInpQ github.com/k0kubun/colorstring v0.0.0-20150214042306-9440f1994b88/go.mod h1:3w7q1U84EfirKl04SVQ/s7nPm1ZPhiXd34z40TNz36k= github.com/k0kubun/pp v3.0.1+incompatible h1:3tqvf7QgUnZ5tXO6pNAZlrvHgl6DvifjDrd9g2S9Z40= github.com/k0kubun/pp v3.0.1+incompatible/go.mod h1:GWse8YhT0p8pT4ir3ZgBbfZild3tgzSScAn6HmfYukg= -github.com/konsorten/go-windows-terminal-sequences v1.0.3 h1:CE8S1cTafDpPvMhIxNJKvHsGVBgn1xWYf1NbHQhywc8= -github.com/konsorten/go-windows-terminal-sequences v1.0.3/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ= github.com/kylelemons/godebug v1.1.0 h1:RPNrshWIDI6G2gRW9EHilWtl7Z6Sb1BR0xunSBf0SNc= github.com/kylelemons/godebug v1.1.0/go.mod h1:9/0rRGxNHcop5bhtWyNeEfOS8JIWk580+fNqagV/RAw= github.com/lucasb-eyer/go-colorful v1.0.2/go.mod h1:0MS4r+7BZKSJ5mw4/S5MPN+qHFF1fYclkSPilDOKW0s= @@ -57,20 +55,16 @@ github.com/nsf/termbox-go v0.0.0-20200204031403-4d2b513ad8be h1:yzmWtPyxEUIKdZg4 github.com/nsf/termbox-go v0.0.0-20200204031403-4d2b513ad8be/go.mod h1:IuKpRQcYE1Tfu+oAQqaLisqDeXgjyyltCfsaoYN18NQ= github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= -github.com/sirupsen/logrus v1.6.0 h1:UBcNElsrwanuuMsnGSlYmtmgbb23qDR5dG+6X6Oo89I= -github.com/sirupsen/logrus v1.6.0/go.mod h1:7uNnSEd1DgxDLC74fIahvMZmmYsHGZGEOFrfsX/uA88= github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA= github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg= github.com/streadway/quantile v0.0.0-20150917103942-b0c588724d25 h1:7z3LSn867ex6VSaahyKadf4WtSsJIgne6A1WLOAGM8A= github.com/streadway/quantile v0.0.0-20150917103942-b0c588724d25/go.mod h1:lbP8tGiBjZ5YWIc2fzuRpTaz0b/53vT6PEs3QuAWzuU= -github.com/stretchr/testify v1.2.2 h1:bSDNvY7ZPG5RlJ8otE/7V6gMiyenm9RtJ7IUVIAoJ1w= -github.com/stretchr/testify v1.2.2/go.mod h1:a8OnRcib4nhh0OaRAV+Yts87kKdq0PP7pXfy6kDkUVs= +github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME= +github.com/stretchr/testify v1.6.1 h1:hDPOHmpOpP40lSULcqw7IrRb/u7w6RpDC9399XyoNd0= +github.com/stretchr/testify v1.6.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg= github.com/tsenart/go-tsz v0.0.0-20180814232043-cdeb9e1e981e/go.mod h1:SWZznP1z5Ki7hDT2ioqiFKEse8K9tU2OUvaRI0NeGQo= github.com/tsenart/vegeta/v12 v12.8.3 h1:UEsDkSrEJojMKW/xr7KUv4H/bYykX+V48KCsPZPqEfk= github.com/tsenart/vegeta/v12 v12.8.3/go.mod h1:ZiJtwLn/9M4fTPdMY7bdbIeyNeFVE8/AHbWFqCsUuho= -github.com/yudai/pp v1.3.0 h1:8qxHWJnYPF9HC9nq1U4N/8toJNBWO4DxaMpDlzrcE6I= -github.com/yudai/pp v2.0.1+incompatible h1:Q4//iY4pNF6yPLZIigmvcl7k/bPgrcTPIFIcmawg5bI= -github.com/yudai/pp v2.0.1+incompatible/go.mod h1:PuxR/8QJ7cyCkFp/aUDS+JY727OFEZkTdatxwunjIkc= golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w= golang.org/x/crypto v0.0.0-20190829043050-9756ffdc2472/go.mod h1:yigFU9vqHzYiE8UmvKecakEJjdnWj3jj499lnFckfCI= golang.org/x/net v0.0.0-20190404232315-eb5bcb51f2a3/go.mod h1:t9HGtf8HONx5eT2rtn7q6eTqICYqUVnKs3thJo3Qplg= @@ -80,7 +74,6 @@ golang.org/x/net v0.0.0-20190827160401-ba9fcec4b297/go.mod h1:z5CRVTTTmAJ677TzLL golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM= golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY= golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= -golang.org/x/sys v0.0.0-20190422165155-953cdadca894/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20190626150813-e07cf5db2756/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= golang.org/x/sys v0.0.0-20190904154756-749cb33beabd h1:DBH9mDw0zluJT/R+nGuV3jWFWLFaHyYZWD4tOT+cjn0= golang.org/x/sys v0.0.0-20190904154756-749cb33beabd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= @@ -93,5 +86,9 @@ golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20190907020128-2ca718005c18/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM= +gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo= +gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= pgregory.net/rapid v0.3.3 h1:jCjBsY4ln4Atz78QoBWxUEvAHaFyNDQg9+WU62aCn1U= pgregory.net/rapid v0.3.3/go.mod h1:UYpPVyjFHzYBGHIxLFoupi8vwk6rXNzRY9OMvVxFIOU= diff --git a/main.go b/main.go index 3c555f2..fd49fb1 100644 --- a/main.go +++ b/main.go @@ -72,6 +72,7 @@ func (c *cli) run() int { func setDebug(w io.Writer, debug bool) { if !debug { pp.SetDefaultOutput(ioutil.Discard) + return } if w == nil { var err error