如何在Golang中为cmd工具编写基准测试？

Question

When I say cmd tool, I mean a tool program like:

var (
    m           = flag.String("m", "GET", "")
    headers     = flag.String("h", "", "")
    body        = flag.String("d", "", "")
    contentType = flag.String("T", "text/html", "")
    ......
)

func main() {
    // Using args above, create complex logics
}

I cannot just call main() in my Benchmark().

Is there any way to do this?

Or must I rewrite a function with args for Benchmark() to call?

Answer 1

It's a good practice to keep your main() as short as possible and delegate all the actual work to a function/method, which makes it testable and benchmark'able if you wish. It may look something along the lines of

// main.go
// go run main.go -h "Content-Type: application/json" -h "X-Request-ID: 12345" -d "Hello"
package main

import (
    "flag"
    "fmt"
    "io"
    "log"
    "strings"
)

type headers map[string]string

func (h headers) String() string {
    hdr := make([]string, 0)
    for k, v := range map[string]string(h) {
        hdr = append(hdr, fmt.Sprintf("%s: %s", k, v))
    }
    return strings.Join(hdr, ",")
}
func (h headers) Set(value string) error {
    i := strings.Index(value, ":")
    if i < 1 {
        return fmt.Errorf("invalid header '%s'", value)
    }
    h[value[0:i]] = strings.TrimSpace(value[i+1:])
    return nil
}

var (
    method string
    hdrs   = make(headers)
    body   string
)

func init() {
    flag.StringVar(&method, "m", "GET", "Method")
    flag.Var(hdrs, "h", "Headers")
    flag.StringVar(&body, "d", "", "Body")
    flag.Parse()
}

func main() {
    err := run(method, hdrs, strings.NewReader(body))
    if err != nil {
        log.Fatal(err)
    }
}

func run(method string, headers headers, body io.Reader) error {
    // implement your logic here
    return nil
}

and your benchmark

// main_test.go
package main

import (
    "strings"
    "testing"
)

func BenchmarkRun1(b *testing.B) {
    hdrs := headers{"Content-Type": "application/json", "Accept": "appication/json"}
    body := "Hello World"
    for i := 0; i < b.N; i++ {
        run("GET", hdrs, strings.NewReader(body))
    }
}