vegeta script

config.ini

[address]
null = "POST http://172.17.1.36:30898/hello"
sleep = "POST http://172.17.1.36:30898/sleep"
query = "POST http://172.17.1.36:30898/es"
[param]
rate = 2000

main.go

package main

import (
	"fmt"
	"github.com/go-ini/ini"
	"os"
	"os/exec"
)

func main() {
	run()
}

func run() {
	addr, _, param, err := getIni()
	if err != nil {
		fmt.Printf("无法加载配置文件: %v", err)
		return
	}

	rate := param.Key("rate").MustInt(20)
	// 执行命令
	nullAddr := addr.Key("null").String()
	sleepAddr := addr.Key("sleep").String()
	queryAddr := addr.Key("query").String()
	invokeAddrs := map[string]string{}
	invokeAddrs["null"] = nullAddr
	invokeAddrs["sleep"] = sleepAddr
	invokeAddrs["query"] = queryAddr
	err = os.Mkdir("result", os.ModePerm)
	if err != nil {
		if os.IsExist(err) {
			// fmt.Printf("目录 %s 已存在\n", "result")
			// 如果目录已存在，可以选择不做任何操作或者继续执行程序
			// 继续执行程序的话，可以在这里添加相应的逻辑
		} else {
			// 如果发生其他错误，可以在这里处理
			fmt.Println("创建目录时发生错误:", err)
		}
	}
	for k, v := range invokeAddrs {
		fmt.Println("当前测试用例：" + k)
		vegeta(k, v, rate)
		fmt.Println("============================")
	}
}

func vegeta(name string, addr string, rate int) {

	// 1. 调用地址生成bin文件
	// 2. 生成report.txt
	// 3. 生成html
	err := generateBinFile(name, addr, rate)
	if err != nil {
		fmt.Printf("无法生成bin文件: %v\n", err)
		return
	}

	// 2. 生成report.txt
	err = generateReport(name)
	if err != nil {
		fmt.Printf("无法生成report文件: %v\n", err)
		return
	}

	// 3. 生成html
	err = generateHTML(name)
	if err != nil {
		fmt.Printf("无法生成html文件: %v\n", err)
		return
	}
}

func generateBinFile(name string, addr string, rate int) error {
	// Code to generate bin file from address
	arg1 := fmt.Sprintf("echo \"%s\" | vegeta attack -rate=%d -duration=1s -body=./query.json -timeout=0 -name=%s | tee ./result/results-%s.bin | vegeta report", addr, rate, name, name)
	_, err := exec.Command("sh", "-c", arg1).Output()
	if err != nil {
		return err
	}
	fmt.Println("生成bin文件成功!")
	return nil
}

func generateReport(name string) error {
	// Code to generate report.txt
	arg := fmt.Sprintf("vegeta report -type=text ./result/results-%s.bin > ./result/report-%s.txt", name, name)
	_, err := exec.Command("sh", "-c", arg).Output()
	if err != nil {
		return err
	}
	fmt.Println("生成report文件成功!")
	return nil
}

func generateHTML(name string) error {
	// Code to generate HTML
	arg := fmt.Sprintf("vegeta plot ./result/results-%s.bin > ./result/plot-%s.html", name, name)
	_, err := exec.Command("sh", "-c", arg).Output()
	if err != nil {
		return err
	}
	fmt.Println("生成HTML文件成功!")
	return nil
}

func getIni() (address *ini.Section, file *ini.Section, param *ini.Section, result error) {
	// 读取配置文件
	cfg, err := ini.Load("config.ini")
	if err != nil {
		fmt.Printf("无法加载配置文件: %v", err)
		return nil, nil, nil, err
	}

	// 获取某个 section 下的键值对
	address, err = cfg.GetSection("address")
	if err != nil {
		return nil, nil, nil, err
	}

	param, err = cfg.GetSection("param")
	if err != nil {
		return nil, nil, nil, err
	}
	return address, nil, param, nil
}

query.json

{
  "start": "2023-12-01 18:05:27",
  "end": "2023-12-01 19:15:27"
}