Merge pull request #607 from justinabrahms/master

Support for specifying additional arguments to concord tasks
2026-04-15 06:57:34 +00:00 · 2020-06-16 09:26:37 +03:00
parent f2eca79a1f dcab2d518f
commit 6f65f6096d
2 changed files with 187 additions and 5 deletions
@@ -13,6 +13,7 @@ import (
 	"net/url"
 	"os"
 	"strconv"
+	"strings"
 	"time"

 	"go.uber.org/zap"
@@ -30,6 +31,8 @@ const defaultPollTimeout = 30
 const concordStatusSuccess = "FINISHED"
 const concordStatusFailed = "FAILED"

+const ARGUMENTS_INPUT_PREFIX = "arguments."
+
 // ConcordTask represents a concord task
 type ConcordTask struct {
 	TaskBase
@@ -38,6 +41,7 @@ type ConcordTask struct {
 	Project      string
 	Repo         string
 	Entrypoint   string
+	Arguments    map[string]string
 	APIKeyPath   string
 	Endpoint     string
 	PollInterval time.Duration
@@ -49,6 +53,7 @@ type ConcordTask struct {
 // NewConcordTask instantiates a new Concord Task
 func NewConcordTask(metadata map[string]string, canary string, logger *zap.SugaredLogger) (*ConcordTask, error) {
 	var pollIntervalInt, pollTimeoutInt int
+	var arguments = make(map[string]string)

 	if _, found := metadata["server"]; !found {
 		return nil, errors.New("`server` is required with type concord")
@@ -96,6 +101,16 @@ func NewConcordTask(metadata map[string]string, canary string, logger *zap.Sugar
 		}
 	}

+	for key, value := range metadata {
+		if key == "arguments.endpoint" {
+			return nil, errors.New("You cannot override Endpoint through arguments. You must override Endpoint directly")
+		}
+
+		if strings.HasPrefix(key, ARGUMENTS_INPUT_PREFIX) {
+			arguments[key[len(ARGUMENTS_INPUT_PREFIX):]] = value
+		}
+	}
+
 	return &ConcordTask{
 		TaskBase: TaskBase{
 			logger: logger,
@@ -107,6 +122,7 @@ func NewConcordTask(metadata map[string]string, canary string, logger *zap.Sugar
 		Entrypoint:   metadata["entrypoint"],
 		APIKeyPath:   metadata["apiKeyPath"],
 		Endpoint:     metadata["endpoint"],
+		Arguments:    arguments,
 		PollInterval: time.Duration(pollIntervalInt) * time.Second,
 		PollTimeout:  time.Duration(pollTimeoutInt) * time.Second,
 		httpClient:   &http.Client{Timeout: 60 * time.Second},
@@ -181,14 +197,22 @@ func (task *ConcordTask) do(req *http.Request, v interface{}) (*http.Response, e
 	return resp, err
 }

+func (task *ConcordTask) buildFields(w *multipart.Writer) {
+	_ = w.WriteField("org", task.Org)
+	_ = w.WriteField("project", task.Project)
+	_ = w.WriteField("repo", task.Repo)
+	_ = w.WriteField("entryPoint", task.Entrypoint)
+	_ = w.WriteField("arguments.endpoint", task.Endpoint)
+
+	for key, value := range task.Arguments {
+		_ = w.WriteField(fmt.Sprintf("arguments.%s", key), value)
+	}
+}
+
 func (task *ConcordTask) startProcess() (string, error) {
 	requestBody := new(bytes.Buffer)
 	writer := multipart.NewWriter(requestBody)
-	_ = writer.WriteField("org", task.Org)
-	_ = writer.WriteField("project", task.Project)
-	_ = writer.WriteField("repo", task.Repo)
-	_ = writer.WriteField("entryPoint", task.Entrypoint)
-	_ = writer.WriteField("arguments.endpoint", task.Endpoint)
+	task.buildFields(writer)

 	err := writer.Close()
 	if err != nil {
@@ -0,0 +1,158 @@
+package loadtester
+
+import (
+	"bytes"
+	"fmt"
+	"io/ioutil"
+	"mime/multipart"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"go.uber.org/zap"
+)
+
+func TestNewConcordTask_Successful(t *testing.T) {
+	metadata := map[string]string{
+		"server":     "example.org",
+		"project":    "my project",
+		"repo":       "my repo",
+		"org":        "my org",
+		"entrypoint": "my-entrypoint",
+		"endpoint":   "example.org/path/to/thing",
+		"apiKeyPath": "/",
+	}
+	task, err := NewConcordTask(metadata, "canary", zap.NewExample().Sugar())
+
+	assert.IsType(t, &ConcordTask{}, task, "Expected to get a well-formed concord task out")
+	assert.Equal(t, nil, err)
+}
+
+func TestNewConcordTask_InitializationWithoutAdequateArgs(t *testing.T) {
+	metadata := map[string]string{
+		"test": "foo",
+	}
+	_, err := NewConcordTask(metadata, "canary", zap.NewExample().Sugar())
+
+	assert.Error(t, err, "is required with type concord")
+}
+
+func TestNewConcordTask_AdditionalArguments(t *testing.T) {
+	metadata := map[string]string{
+		"server":         "example.org",
+		"project":        "my project",
+		"repo":           "my repo",
+		"org":            "my org",
+		"entrypoint":     "my-entrypoint",
+		"endpoint":       "example.org/path/to/thing",
+		"apiKeyPath":     "/",
+		"arguments.test": "works",
+	}
+	task, err := NewConcordTask(metadata, "canary", zap.NewExample().Sugar())
+
+	assert.IsType(t, &ConcordTask{}, task, "Expected to get a well-formed concord task out")
+	assert.Equal(t, nil, err)
+	assert.Equal(t, task.Arguments["test"], "works")
+}
+
+func TestNewConcordTask_DontOverrideEndpoint(t *testing.T) {
+	metadata := map[string]string{
+		"server":             "example.org",
+		"project":            "my project",
+		"repo":               "my repo",
+		"org":                "my org",
+		"entrypoint":         "my-entrypoint",
+		"endpoint":           "example.org/path/to/thing",
+		"apiKeyPath":         "/",
+		"arguments.endpoint": "works",
+	}
+	_, err := NewConcordTask(metadata, "canary", zap.NewExample().Sugar())
+
+	assert.Error(t, err, "You cannot override Endpoint through arguments")
+
+}
+
+func assertNextPartHasKeyAndValue(t *testing.T, r *multipart.Reader, key string, value string) {
+	part, err := r.NextPart()
+	if err != nil {
+		t.Fatalf("Part failed: %v", err)
+	}
+	// assert.Equal(t, 1, part)
+
+	slurp, err := ioutil.ReadAll(part)
+	if err != nil {
+		fmt.Printf("Part: %+v", part)
+		t.Fatalf("Couldn't read part: %v", err)
+	}
+
+	assert.Equal(t, part.FormName(), key)
+	assert.Equal(t, string(slurp), value)
+
+}
+
+func TestConcordTask_BuildingDefaultFields(t *testing.T) {
+	var b bytes.Buffer
+	w := multipart.NewWriter(&b)
+
+	task := &ConcordTask{
+		Org:        "my org",
+		Project:    "my project",
+		Repo:       "my repo",
+		Entrypoint: "my entrypoint",
+		Endpoint:   "example.org",
+	}
+
+	task.buildFields(w)
+
+	err := w.Close()
+	if err != nil {
+		t.Fatalf("Couldn't close writer: %v", err)
+	}
+
+	r := multipart.NewReader(&b, w.Boundary())
+	assertNextPartHasKeyAndValue(t, r, "org", "my org")
+	assertNextPartHasKeyAndValue(t, r, "project", "my project")
+	assertNextPartHasKeyAndValue(t, r, "repo", "my repo")
+	assertNextPartHasKeyAndValue(t, r, "entryPoint", "my entrypoint")
+	assertNextPartHasKeyAndValue(t, r, "arguments.endpoint", "example.org")
+
+	part, _ := r.NextPart()
+	if part != nil {
+		t.Errorf("Didn't expect additional parts, but got %v", part)
+	}
+}
+
+func TestConcordTask_AdditionalArguments(t *testing.T) {
+	var b bytes.Buffer
+	w := multipart.NewWriter(&b)
+
+	task := &ConcordTask{
+		Org:        "my org",
+		Project:    "my project",
+		Repo:       "my repo",
+		Entrypoint: "my entrypoint",
+		Endpoint:   "example.org",
+		Arguments: map[string]string{
+			"test": "thing",
+		},
+	}
+
+	task.buildFields(w)
+
+	err := w.Close()
+	if err != nil {
+		t.Fatalf("Couldn't close writer: %v", err)
+	}
+
+	r := multipart.NewReader(&b, w.Boundary())
+	assertNextPartHasKeyAndValue(t, r, "org", "my org")
+	assertNextPartHasKeyAndValue(t, r, "project", "my project")
+	assertNextPartHasKeyAndValue(t, r, "repo", "my repo")
+	assertNextPartHasKeyAndValue(t, r, "entryPoint", "my entrypoint")
+	assertNextPartHasKeyAndValue(t, r, "arguments.endpoint", "example.org")
+	assertNextPartHasKeyAndValue(t, r, "arguments.test", "thing")
+
+	part, _ := r.NextPart()
+	if part != nil {
+		t.Errorf("Didn't expect additional parts, but got %v", part)
+	}
+}