oef/benchmarks/gen_targets.go

package main

import (
	"bytes"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path"
	"path/filepath"
	"strings"
	"text/template"

	"git.andreafazzi.eu/andrea/oef/client"
	"git.andreafazzi.eu/andrea/oef/orm"
	"github.com/gocarina/gocsv"
	vegeta "github.com/tsenart/vegeta/lib"
)

type config struct {
	Url               string
	SubscribePassword string
}

type genFunc func(string, *config) error

var (
	generators  map[string]genFunc
	templateDir string = "./templates"
)

func incr(value int) int {
	return value + 1
}

func genGETschoolSubscription(targetName string, conf *config) error {
	token, err := client.GetToken(conf.Url, "", conf.SubscribePassword)
	if err != nil {
		return err
	}

	err = os.Mkdir("targets", 0777)
	if errors.Is(err, &os.PathError{}) {
		return err
	}

	tmplPath := path.Join(templateDir, targetName+".tpl")

	log.Printf("Parse %s template", tmplPath)
	tmpl, err := template.New(filepath.Base(tmplPath)).Funcs(template.FuncMap{"incr": incr}).ParseFiles(tmplPath)
	if err != nil {
		return err
	}

	output, err := os.Create("./targets/school_subscription.txt")
	if err != nil {
		return err
	}
	defer output.Close()

	log.Printf("Execute %s template and generate targets...", targetName)
	data := struct {
		Token  string
		Config *config
	}{token, conf}

	err = tmpl.Execute(output, data)
	if err != nil {
		return err
	}

	return nil
}

func genPOSTResponseTargets(targetName string, conf *config) error {
	var (
		participants []*orm.Participant
		tokens       map[string]string
	)

	log.Println("Read participants.csv...")

	input, err := ioutil.ReadFile("./testdata/participants.csv")
	if err != nil {
		return err
	}
	if err := gocsv.Unmarshal(strings.NewReader(string(input)), &participants); err != nil {
		panic(err)
	}

	log.Println("Obtaining tokens for each participants...")
	tokens = make(map[string]string, 0)
	for _, participant := range participants {
		token, err := client.GetToken(conf.Url, participant.FiscalCode, participant.Password)
		if err != nil {
			return err
		}
		tokens[participant.FiscalCode] = token
	}

	targets := make([]*vegeta.Target, 0)
	form := url.Values{}
	form.Set("Singleresponses.0", "32")

	for i, participant := range participants {
		targets = append(targets, &vegeta.Target{
			Method: "POST",
			URL:    fmt.Sprintf("%s/responses/%d/update?format=html&tpl_content=responses_add_update&tpl_layout=base&update=true&login_session=%s", conf.Url, i+1, tokens[participant.FiscalCode]),
			Body:   []byte(form.Encode()),
			Header: http.Header{"Content-Type": []string{"application/x-www-form-urlencoded"}},
		})
	}

	err = os.Mkdir("targets", 0777)
	if errors.Is(err, &os.PathError{}) {
		return err
	}
	output, err := os.Create("./targets/post_response_targets.txt")
	if err != nil {
		return err
	}
	defer output.Close()

	for _, t := range targets {
		var line string
		buf := bytes.NewBufferString(line)
		encoder := json.NewEncoder(buf)
		encoder.SetEscapeHTML(false)
		err := encoder.Encode(t)
		// jsonData, err := json.Marshal(t)
		if err != nil {
			return err
		}
		_, err = output.WriteString(buf.String() + "\n\n")
		if err != nil {
			return err
		}
	}

	return nil
}

func genGETResponseTargets(targetName string, conf *config) error {
	var data struct {
		Participants []*orm.Participant
		Tokens       map[string]string
		Config       *config
	}

	data.Config = conf

	log.Println("Read participants.csv...")

	input, err := ioutil.ReadFile("./testdata/participants.csv")
	if err != nil {
		return err
	}
	if err := gocsv.Unmarshal(strings.NewReader(string(input)), &data.Participants); err != nil {
		panic(err)
	}

	log.Println("Obtaining tokens for each participants...")
	data.Tokens = make(map[string]string, 0)
	for _, participant := range data.Participants {
		token, err := client.GetToken(data.Config.Url, participant.FiscalCode, participant.Password)
		if err != nil {
			return err
		}
		data.Tokens[participant.FiscalCode] = token
	}

	tmplPath := path.Join(templateDir, targetName+".tpl")

	log.Printf("Parse %s template", tmplPath)
	tmpl, err := template.New(filepath.Base(tmplPath)).Funcs(template.FuncMap{"incr": incr}).ParseFiles(tmplPath)
	if err != nil {
		return err
	}

	err = os.Mkdir("targets", 0777)
	if errors.Is(err, &os.PathError{}) {
		return err
	}
	output, err := os.Create("./targets/response_targets.txt")
	if err != nil {
		return err
	}
	defer output.Close()

	log.Printf("Execute %s template and generate targets...", targetName)
	err = tmpl.Execute(output, data)
	if err != nil {
		return err
	}

	return nil
}

func init() {
	generators = map[string]genFunc{
		"responses":           genGETResponseTargets,
		"post_responses":      genPOSTResponseTargets,
		"school_subscription": genGETschoolSubscription,
	}

}

func main() {
	target := flag.String("targets", "responses", "Generate targets for participants responses.")
	url := flag.String("url", "http://localhost:3000", "The URL of the host.")
	flag.Parse()

	fn, ok := generators[*target]
	if !ok {
		log.Fatal(errors.New("Unknown target"))
	}

	err := fn(*target, &config{Url: *url, SubscribePassword: "subscribe"})
	if err != nil {
		panic(err)
	}

}