1 files changed, 36 insertions, 6 deletions
diff --git a/example/example.go b/example/example.go
index 44c3b5b..33c30d1 100644
--- a/example/example.go
+++ b/example/example.go
@@ -10,8 +10,10 @@ import (
 	"net/http"
 	"time"
 
-	"git.sr.ht/~shulhan/trunks"
 	libhttp "github.com/shuLhan/share/lib/http"
+	vegeta "github.com/tsenart/vegeta/v12/lib"
+
+	"git.sr.ht/~shulhan/trunks"
 )
 
 const (
@@ -19,7 +21,8 @@ const (
 )
 
 type Example struct {
-	trunks *trunks.Trunks
+	trunks           *trunks.Trunks
+	targetExampleGet vegeta.Target
 }
 
 //
@@ -27,7 +30,7 @@ type Example struct {
 //
 func New() (ex *Example, err error) {
 	env := &trunks.Environment{
-		ResultsDir:    "/tmp",
+		ResultsDir:    "testdata/example/",
 		ResultsSuffix: "_trunks_example",
 	}
 
@@ -76,7 +79,7 @@ func (ex *Example) registerEndpoints() (err error) {
 func (ex *Example) registerTargets() (err error) {
 	targetHttp := &trunks.Target{
 		Name: "Example HTTP target",
-		AttackOpts: &trunks.AttackOptions{
+		Opts: &trunks.AttackOptions{
 			BaseUrl:       fmt.Sprintf("http://%s", ex.trunks.Env.ListenAddress),
 			Duration:      5 * time.Second,
 			RatePerSecond: 10,
@@ -92,7 +95,10 @@ func (ex *Example) registerTargets() (err error) {
 			Params: trunks.KeyValue{
 				"Param1": "1",
 			},
-			Run: ex.runExampleGet,
+			Run:         ex.runExampleGet,
+			Attack:      ex.attackExampleGet,
+			PreAttack:   ex.preattackExampleGet,
+			AllowAttack: true,
 		}},
 	}
 
@@ -112,7 +118,7 @@ func (ex *Example) pathExampleGet(epr *libhttp.EndpointRequest) ([]byte, error)
 
 func (ex *Example) runExampleGet(target *trunks.Target, req *trunks.RunRequest) ([]byte, error) {
 	if target.HttpClient == nil {
-		target.HttpClient = libhttp.NewClient(target.AttackOpts.BaseUrl, nil, true)
+		target.HttpClient = libhttp.NewClient(target.Opts.BaseUrl, nil, true)
 	}
 	_, resbody, err := target.HttpClient.Get(
 		req.HttpTarget.Path,
@@ -123,3 +129,27 @@ func (ex *Example) runExampleGet(target *trunks.Target, req *trunks.RunRequest)
 	}
 	return resbody, nil
 }
+
+func (ex *Example) preattackExampleGet(rr *trunks.RunRequest) {
+	ex.targetExampleGet = vegeta.Target{
+		Method: rr.HttpTarget.Method.String(),
+		URL:    fmt.Sprintf("%s%s", rr.Target.Opts.BaseUrl, rr.HttpTarget.Path),
+		Header: rr.HttpTarget.Headers.ToHttpHeader(),
+	}
+
+	q := rr.HttpTarget.Params.ToUrlValues().Encode()
+	if len(q) > 0 {
+		ex.targetExampleGet.URL += "?" + q
+	}
+
+	fmt.Printf("preattackExampleGet: %+v\n", ex.targetExampleGet)
+}
+
+func (ex *Example) attackExampleGet(rr *trunks.RunRequest) vegeta.Targeter {
+	return func(tgt *vegeta.Target) error {
+		rr.HttpTarget.AttackLocker.Lock()
+		*tgt = ex.targetExampleGet
+		rr.HttpTarget.AttackLocker.Unlock()
+		return nil
+	}
+}