Files
load-testing/main.go
2021-07-18 14:07:08 -06:00

139 lines
2.7 KiB
Go

package main
import (
"fmt"
"log"
"net/http"
"os"
"time"
"github.com/urfave/cli"
)
var app *cli.App
func main() {
app = &cli.App{
Flags: []cli.Flag{
&cli.StringFlag{
Name: "method",
Aliases: []string{"m"},
Value: "GET",
Usage: "The `METHOD` for making the request",
},
&cli.DurationFlag{
Name: "duration",
Aliases: []string{"d"},
Value: time.Second * 60,
Usage: "`DURATION` for how long to run the test",
},
&cli.Int64Flag{
Name: "rpm",
Usage: "`RATE-PER-MINUTE`",
Value: 60,
},
&cli.StringSliceFlag{
Name: "request",
Aliases: []string{"r"},
Value: cli.NewStringSlice("http://localhost:8080/"),
Usage: "The `URL` for making the request",
},
&cli.StringSliceFlag{
Name: "header",
Aliases: []string{"H"},
Value: cli.NewStringSlice("Content-Type: application/json"),
Usage: "The `HEADER` for the request",
},
},
Name: "Burst CLI",
Usage: "Run Burst to throw large amounts of traffic at a give url or set of urls",
UsageText: "burst -m GET -r http://localhost:8080/ -H 'Content-Type: application/json' -o data.json --rpm 1000",
Action: func(c *cli.Context) error {
fmt.Println("Bursting with happiness")
process(c)
return nil
},
EnableBashCompletion: true,
Version: "1.0.0",
}
err := app.Run(os.Args)
if err != nil {
log.Fatal(err)
}
}
func request(ctx *cli.Context, i int) (*http.Response, error) {
req := ctx.StringSlice("request")
var resp *http.Response
method := ctx.String("method")
if method == "GET" {
resp, err := http.Get(req[i])
if err != nil {
log.Fatal(err)
}
fmt.Println(resp.Body)
}
// if method == "POST" {
// resp, err := http.POST(req[i])
// if err != nil {
// log.Fatal(err)
// }
// fmt.Println(resp.Body)
// }
return resp, nil
}
func process(ctx *cli.Context) error {
req := ctx.StringSlice("request")
ticker := time.NewTicker(500 * time.Millisecond)
done := make(chan bool, len(req))
go func() {
for {
for {
select {
case <-done:
return
case t := <-ticker.C:
fmt.Println(req[0])
resp, err := request(ctx, 0)
if err != nil {
log.Fatal(err)
}
fmt.Println(resp.Body)
fmt.Println("Tick at", t)
}
}
}
}()
time.Sleep(ctx.Duration("duration"))
done <- true
// select {
// case <-ch:
// fmt.Println("Read from ch")
// case <-time.After(1 * time.Second):
// fmt.Println("Timed out")
// }
// if method == "POST" {
// resp, err := http.POST(req, data)
// if err != nil {
// log.Fatal(err)
// return err
// }
// }
// if method == "PUT" {
// resp, err := http.Get(req)
// if err != nil {
// log.Fatal(err)
// return err
// }
// }
return nil
}