Modularize flag parsing

main.go

@@ -49,32 +49,39 @@ type cli struct {
 }
 
 func main() {
+ c, err := parseFlags(os.Stdout, os.Stderr)
+ if err != nil {
+ os.Exit(0)
+ }
+ os.Exit(c.run(flagSet.Args()))
+}
+
+func parseFlags(stdout, stderr io.Writer) (*cli, error) {
  c := &cli{
- stdout: os.Stdout,
- stderr: os.Stderr,
+ stdout: stdout,
+ stderr: stderr,
  }
- flagSet.IntVarP(&c.rate, "rate", "r", 50, "The request rate per second to issue against the targets. Give 0 then it will send requests as fast as possible.")
- flagSet.DurationVarP(&c.duration, "duration", "d", time.Second*10, "The amount of time to issue requests to the targets. Give 0s for an infinite attack.")
- flagSet.DurationVarP(&c.timeout, "timeout", "t", time.Second*30, "The timeout for each request. 0s means to disable timeouts.")
- flagSet.StringVarP(&c.method, "method", "m", "GET", "An HTTP request method for each request.")
+ flagSet.IntVarP(&c.rate, "rate", "r", attacker.DefaultRate, "The request rate per second to issue against the targets. Give 0 then it will send requests as fast as possible.")
+ flagSet.DurationVarP(&c.duration, "duration", "d", attacker.DefaultDuration, "The amount of time to issue requests to the targets. Give 0s for an infinite attack.")
+ flagSet.DurationVarP(&c.timeout, "timeout", "t", attacker.DefaultTimeout, "The timeout for each request. 0s means to disable timeouts.")
+ flagSet.StringVarP(&c.method, "method", "m", attacker.DefaultMethod, "An HTTP request method for each request.")
  flagSet.StringSliceVarP(&c.headers, "header", "H", []string{}, "A request header to be sent. Can be used multiple times to send multiple headers.")
  flagSet.StringVarP(&c.body, "body", "b", "", "A request body to be sent.")
  flagSet.StringVarP(&c.bodyFile, "body-file", "B", "", "The path to file whose content will be set as the http request body.")
- flagSet.Int64VarP(&c.maxBody, "max-body", "M", -1, "Max bytes to capture from response bodies. Give -1 for no limit.")
+ flagSet.Int64VarP(&c.maxBody, "max-body", "M", attacker.DefaultMaxBody, "Max bytes to capture from response bodies. Give -1 for no limit.")
  flagSet.BoolVarP(&c.version, "version", "v", false, "Print the current version.")
  flagSet.BoolVar(&c.debug, "debug", false, "Run in debug mode.")
  flagSet.BoolVarP(&c.keepAlive, "keepalive", "k", true, "Use persistent connections.")
- flagSet.Uint64VarP(&c.workers, "workers", "w", attacker.DefaultWorkers, "Amount of workers to spawn.")
+ flagSet.Uint64VarP(&c.workers, "workers", "w", attacker.DefaultWorkers, "Amount of initial workers to spawn.")
  flagSet.Uint64VarP(&c.maxWorkers, "max-workers", "W", attacker.DefaultMaxWorkers, "Amount of maximum workers to spawn.")
  flagSet.Usage = c.usage
  if err := flagSet.Parse(os.Args[1:]); err != nil {
  if !errors.Is(err, flag.ErrHelp) {
  fmt.Fprintln(c.stderr, err)
  }
- return
+ return nil, err
  }
-
- os.Exit(c.run(flagSet.Args()))
+ return c, nil
 }
 
 func (c *cli) run(args []string) int {

main_test.go

@@ -2,12 +2,14 @@ package main
 
 import (
  "bytes"
+ "math"
  "net/http"
  "testing"
-
- "github.com/nakabonne/ali/attacker"
+ "time"
 
  "github.com/stretchr/testify/assert"
+
+ "github.com/nakabonne/ali/attacker"
 )
 
 func TestValidateMethod(t *testing.T) {
@@ -35,6 +37,40 @@ func TestValidateMethod(t *testing.T) {
  }
 }
 
+func TestParseFlags(t *testing.T) {
+ tests := []struct {
+ name string
+ want *cli
+ wantErr bool
+ }{
+ {
+ name: "with default options",
+ want: &cli{
+ rate: 50,
+ duration: time.Second * 10,
+ timeout: time.Second * 30,
+ method: "GET",
+ headers: []string{},
+ maxBody: -1,
+ keepAlive: true,
+ workers: 10,
+ maxWorkers: math.MaxUint64,
+ stdout: new(bytes.Buffer),
+ stderr: new(bytes.Buffer),
+ },
+ wantErr: false,
+ },
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ b := new(bytes.Buffer)
+ got, err := parseFlags(b, b)
+ assert.Equal(t, tt.want, got)
+ assert.Equal(t, tt.wantErr, err != nil)
+ })
+ }
+}
+
 func TestRun(t *testing.T) {
  tests := []struct {
  name string
@@ -152,6 +188,16 @@ func TestMakeOptions(t *testing.T) {
  },
  wantErr: false,
  },
+ {
+ name: "wrong body file given",
+ cli: &cli{
+ method: "GET",
+ headers: []string{"key:value"},
+ bodyFile: "wrong",
+ },
+ want: nil,
+ wantErr: true,
+ },
  }
  for _, tt := range tests {
  t.Run(tt.name, func(t *testing.T) {
@@ -161,3 +207,24 @@ func TestMakeOptions(t *testing.T) {
  })
  }
 }
+
+func TestSetDebug(t *testing.T) {
+ tests := []struct {
+ name string
+ debug bool
+ }{
+ {
+ name: "in non-debug use",
+ debug: false,
+ },
+ {
+ name: "in debug use",
+ debug: true,
+ },
+ }
+ for _, tt := range tests {
+ t.Run(tt.name, func(t *testing.T) {
+ setDebug(nil, tt.debug)
+ })
+ }
+}