loadtester: return cmd output optionally

pkg/loadtester/bash.go

@@ -18,20 +18,20 @@ func (task *BashTask) Hash() string {
 	return hash(task.canary + task.command)
 }
 
-func (task *BashTask) Run(ctx context.Context) (bool, error) {
+func (task *BashTask) Run(ctx context.Context) (*TaskRunResult, error) {
 	cmd := exec.CommandContext(ctx, "bash", "-c", task.command)
 	out, err := cmd.CombinedOutput()
 
 	if err != nil {
 		task.logger.With("canary", task.canary).Errorf("command failed %s %v %s", task.command, err, out)
-		return false, fmt.Errorf("command %s failed: %s: %w", task.command, out, err)
+		return &TaskRunResult{false, out}, fmt.Errorf("command %s failed: %s: %w", task.command, out, err)
 	} else {
 		if task.logCmdOutput {
 			fmt.Printf("%s\n", out)
 		}
 		task.logger.With("canary", task.canary).Infof("command finished %s", task.command)
 	}
-	return true, nil
+	return &TaskRunResult{true, out}, nil
 }
 
 func (task *BashTask) String() string {

pkg/loadtester/concord.go

@@ -121,14 +121,15 @@ func (task *ConcordTask) String() string {
 	return fmt.Sprintf("%s %s %s %s", task.Org, task.Project, task.Repo, task.Entrypoint)
 }
 
-func (task *ConcordTask) Run(ctx context.Context) (bool, error) {
+func (task *ConcordTask) Run(ctx context.Context) (*TaskRunResult, error) {
 	instance, err := task.startProcess()
 	if err != nil {
 		task.logger.Errorf("failed to start process: %s", err.Error())
-		return false, err
+		return &TaskRunResult{false, nil}, err
 	}
 
-	return task.checkStatus(ctx, instance, task.PollInterval)
+	ok, err := task.checkStatus(ctx, instance, task.PollInterval)
+	return &TaskRunResult{ok, nil}, err
 }
 
 type concordProcess struct {

pkg/loadtester/helm.go

@@ -19,7 +19,7 @@ func (task *HelmTask) Hash() string {
 	return hash(task.canary + task.command)
 }
 
-func (task *HelmTask) Run(ctx context.Context) (bool, error) {
+func (task *HelmTask) Run(ctx context.Context) (*TaskRunResult, error) {
 	helmCmd := fmt.Sprintf("%s %s", TaskTypeHelm, task.command)
 	task.logger.With("canary", task.canary).Infof("running command %v", helmCmd)
 
@@ -27,14 +27,14 @@ func (task *HelmTask) Run(ctx context.Context) (bool, error) {
 	out, err := cmd.CombinedOutput()
 	if err != nil {
 		task.logger.With("canary", task.canary).Errorf("command failed %s %v %s", task.command, err, out)
-		return false, fmt.Errorf("command %s failed: %s: %w", task.command, out, err)
+		return &TaskRunResult{false, out}, fmt.Errorf("command %s failed: %s: %w", task.command, out, err)
 	} else {
 		if task.logCmdOutput {
 			fmt.Printf("%s\n", out)
 		}
 		task.logger.With("canary", task.canary).Infof("command finished %v", helmCmd)
 	}
-	return true, nil
+	return &TaskRunResult{true, out}, nil
 }
 
 func (task *HelmTask) String() string {

pkg/loadtester/helmv3.go

@@ -19,7 +19,7 @@ func (task *HelmTaskv3) Hash() string {
 	return hash(task.canary + task.command)
 }
 
-func (task *HelmTaskv3) Run(ctx context.Context) (bool, error) {
+func (task *HelmTaskv3) Run(ctx context.Context) (*TaskRunResult, error) {
 	helmCmd := fmt.Sprintf("%s %s", TaskTypeHelmv3, task.command)
 	task.logger.With("canary", task.canary).Infof("running command %v", helmCmd)
 
@@ -27,14 +27,14 @@ func (task *HelmTaskv3) Run(ctx context.Context) (bool, error) {
 	out, err := cmd.CombinedOutput()
 	if err != nil {
 		task.logger.With("canary", task.canary).Errorf("command failed %s %v %s", task.command, err, out)
-		return false, fmt.Errorf("command %s failed: %s: %w", task.command, out, err)
+		return &TaskRunResult{false, out}, fmt.Errorf("command %s failed: %s: %w", task.command, out, err)
 	} else {
 		if task.logCmdOutput {
 			fmt.Printf("%s\n", out)
 		}
 		task.logger.With("canary", task.canary).Infof("command finished %v", helmCmd)
 	}
-	return true, nil
+	return &TaskRunResult{true, out}, nil
 }
 
 func (task *HelmTaskv3) String() string {

pkg/loadtester/runner.go

@@ -9,6 +9,13 @@ import (
 	"go.uber.org/zap"
 )
 
+type TaskRunnerInterface interface {
+	Add(task Task)
+	GetTotalExecs() uint64
+	Start(interval time.Duration, stopCh <-chan struct{})
+	Timeout() time.Duration
+}
+
 type TaskRunner struct {
 	logger       *zap.SugaredLogger
 	timeout      time.Duration
@@ -81,3 +88,7 @@ func (tr *TaskRunner) Start(interval time.Duration, stopCh <-chan struct{}) {
 		}
 	}
 }
+
+func (tr *TaskRunner) Timeout() time.Duration {
+	return tr.timeout
+}

pkg/loadtester/server.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"io/ioutil"
 	"net/http"
+	"strconv"
 	"time"
 
 	"github.com/prometheus/client_golang/prometheus/promhttp"
@@ -17,10 +18,7 @@ import (
 func ListenAndServe(port string, timeout time.Duration, logger *zap.SugaredLogger, taskRunner *TaskRunner, gate *GateStorage, stopCh <-chan struct{}) {
 	mux := http.DefaultServeMux
 	mux.Handle("/metrics", promhttp.Handler())
-	mux.HandleFunc("/healthz", func(w http.ResponseWriter, r *http.Request) {
-		w.WriteHeader(http.StatusOK)
-		w.Write([]byte("OK"))
-	})
+	mux.HandleFunc("/healthz", HandleHealthz)
 	mux.HandleFunc("/gate/approve", func(w http.ResponseWriter, r *http.Request) {
 		w.WriteHeader(http.StatusOK)
 		w.Write([]byte("OK"))
@@ -187,157 +185,7 @@ func ListenAndServe(port string, timeout time.Duration, logger *zap.SugaredLogge
 		logger.Infof("%s rollback closed", canaryName)
 	})
 
-	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
-		body, err := ioutil.ReadAll(r.Body)
-		if err != nil {
-			logger.Error("reading the request body failed", zap.Error(err))
-			w.WriteHeader(http.StatusBadRequest)
-			return
-		}
-		defer r.Body.Close()
-
-		payload := &flaggerv1.CanaryWebhookPayload{}
-		err = json.Unmarshal(body, payload)
-		if err != nil {
-			logger.Error("decoding the request body failed", zap.Error(err))
-			w.WriteHeader(http.StatusBadRequest)
-			return
-		}
-
-		if len(payload.Metadata) > 0 {
-			metadata := payload.Metadata
-			var typ, ok = metadata["type"]
-			if !ok {
-				typ = TaskTypeShell
-			}
-
-			// run bats command (blocking task)
-			if typ == TaskTypeBash {
-				logger.With("canary", payload.Name).Infof("bats command %s", payload.Metadata["cmd"])
-
-				bats := BashTask{
-					command:      payload.Metadata["cmd"],
-					logCmdOutput: true,
-					TaskBase: TaskBase{
-						canary: fmt.Sprintf("%s.%s", payload.Name, payload.Namespace),
-						logger: logger,
-					},
-				}
-
-				ctx, cancel := context.WithTimeout(context.Background(), taskRunner.timeout)
-				defer cancel()
-
-				ok, err := bats.Run(ctx)
-				if !ok {
-					w.WriteHeader(http.StatusInternalServerError)
-					w.Write([]byte(err.Error()))
-					return
-				}
-
-				w.WriteHeader(http.StatusOK)
-				return
-			}
-
-			// run helm command (blocking task)
-			if typ == TaskTypeHelm {
-				helm := HelmTask{
-					command:      payload.Metadata["cmd"],
-					logCmdOutput: true,
-					TaskBase: TaskBase{
-						canary: fmt.Sprintf("%s.%s", payload.Name, payload.Namespace),
-						logger: logger,
-					},
-				}
-
-				ctx, cancel := context.WithTimeout(context.Background(), taskRunner.timeout)
-				defer cancel()
-
-				ok, err := helm.Run(ctx)
-				if !ok {
-					w.WriteHeader(http.StatusInternalServerError)
-					w.Write([]byte(err.Error()))
-					return
-				}
-
-				w.WriteHeader(http.StatusOK)
-				return
-			}
-
-			// run helmv3 command (blocking task)
-			if typ == TaskTypeHelmv3 {
-				helm := HelmTaskv3{
-					command:      payload.Metadata["cmd"],
-					logCmdOutput: true,
-					TaskBase: TaskBase{
-						canary: fmt.Sprintf("%s.%s", payload.Name, payload.Namespace),
-						logger: logger,
-					},
-				}
-
-				ctx, cancel := context.WithTimeout(context.Background(), taskRunner.timeout)
-				defer cancel()
-
-				ok, err := helm.Run(ctx)
-				if !ok {
-					w.WriteHeader(http.StatusInternalServerError)
-					w.Write([]byte(err.Error()))
-					return
-				}
-
-				w.WriteHeader(http.StatusOK)
-				return
-			}
-
-			// run concord job (blocking task)
-			if typ == TaskTypeConcord {
-				concord, err := NewConcordTask(payload.Metadata, fmt.Sprintf("%s.%s", payload.Name, payload.Namespace), logger)
-
-				if err != nil {
-					logger.With("canary", payload.Name).Errorf("concord task init error: %s", err)
-					w.WriteHeader(http.StatusInternalServerError)
-					w.Write([]byte(err.Error()))
-					return
-				}
-
-				ctx, cancel := context.WithTimeout(context.Background(), taskRunner.timeout)
-				defer cancel()
-
-				ok, err := concord.Run(ctx)
-				if !ok {
-					if err != nil {
-						logger.With("canary", payload.Name).Errorf("concord task error: %s", err)
-					}
-					w.WriteHeader(http.StatusInternalServerError)
-					w.Write([]byte(err.Error()))
-					return
-				}
-
-				w.WriteHeader(http.StatusOK)
-				return
-			}
-
-			taskFactory, ok := GetTaskFactory(typ)
-			if !ok {
-				w.WriteHeader(http.StatusBadRequest)
-				w.Write([]byte(fmt.Sprintf("unknown task type %s", typ)))
-				return
-			}
-			canary := fmt.Sprintf("%s.%s", payload.Name, payload.Namespace)
-			task, err := taskFactory(metadata, canary, logger)
-			if err != nil {
-				w.WriteHeader(http.StatusBadRequest)
-				w.Write([]byte(err.Error()))
-				return
-			}
-			taskRunner.Add(task)
-		} else {
-			w.WriteHeader(http.StatusBadRequest)
-			w.Write([]byte("metadata not found in payload"))
-			return
-		}
-
-		w.WriteHeader(http.StatusAccepted)
-	})
+	mux.HandleFunc("/", HandleNewTask(logger, taskRunner))
 	srv := &http.Server{
 		Addr:         ":" + port,
 		Handler:      mux,
@@ -364,3 +212,181 @@ func ListenAndServe(port string, timeout time.Duration, logger *zap.SugaredLogge
 		logger.Info("HTTP server stopped")
 	}
 }
+
+// HandleHealthz handles heath check requests
+func HandleHealthz(w http.ResponseWriter, r *http.Request) {
+	w.WriteHeader(http.StatusOK)
+	w.Write([]byte("OK"))
+}
+
+// HandleNewTask handles task creation requests
+func HandleNewTask(logger *zap.SugaredLogger, taskRunner TaskRunnerInterface) func(w http.ResponseWriter, r *http.Request) {
+	return func(w http.ResponseWriter, r *http.Request) {
+		body, err := ioutil.ReadAll(r.Body)
+		if err != nil {
+			logger.Error("reading the request body failed", zap.Error(err))
+			w.WriteHeader(http.StatusBadRequest)
+			return
+		}
+		defer r.Body.Close()
+
+		payload := &flaggerv1.CanaryWebhookPayload{}
+		err = json.Unmarshal(body, payload)
+		if err != nil {
+			logger.Error("decoding the request body failed", zap.Error(err))
+			w.WriteHeader(http.StatusBadRequest)
+			return
+		}
+
+		if len(payload.Metadata) > 0 {
+			metadata := payload.Metadata
+			var typ, ok = metadata["type"]
+			if !ok {
+				typ = TaskTypeShell
+			}
+
+			rtnCmdOutput := false
+			if rtn, ok := metadata["returnCmdOutput"]; ok {
+				rtnCmdOutput, err = strconv.ParseBool(rtn)
+			}
+
+			// run bats command (blocking task)
+			if typ == TaskTypeBash {
+				logger.With("canary", payload.Name).Infof("bats command %s", payload.Metadata["cmd"])
+
+				bats := BashTask{
+					command:      payload.Metadata["cmd"],
+					logCmdOutput: true,
+					TaskBase: TaskBase{
+						canary: fmt.Sprintf("%s.%s", payload.Name, payload.Namespace),
+						logger: logger,
+					},
+				}
+
+				ctx, cancel := context.WithTimeout(context.Background(), taskRunner.Timeout())
+				defer cancel()
+
+				result, err := bats.Run(ctx)
+				if !result.ok {
+					w.WriteHeader(http.StatusInternalServerError)
+					w.Write([]byte(err.Error()))
+					return
+				}
+
+				w.WriteHeader(http.StatusOK)
+				if rtnCmdOutput {
+					w.Write(result.out)
+				}
+				return
+			}
+
+			// run helm command (blocking task)
+			if typ == TaskTypeHelm {
+				helm := HelmTask{
+					command:      payload.Metadata["cmd"],
+					logCmdOutput: true,
+					TaskBase: TaskBase{
+						canary: fmt.Sprintf("%s.%s", payload.Name, payload.Namespace),
+						logger: logger,
+					},
+				}
+
+				ctx, cancel := context.WithTimeout(context.Background(), taskRunner.Timeout())
+				defer cancel()
+
+				result, err := helm.Run(ctx)
+				if !result.ok {
+					w.WriteHeader(http.StatusInternalServerError)
+					w.Write([]byte(err.Error()))
+					return
+				}
+
+				w.WriteHeader(http.StatusOK)
+				if rtnCmdOutput {
+					w.Write(result.out)
+				}
+				return
+			}
+
+			// run helmv3 command (blocking task)
+			if typ == TaskTypeHelmv3 {
+				helm := HelmTaskv3{
+					command:      payload.Metadata["cmd"],
+					logCmdOutput: true,
+					TaskBase: TaskBase{
+						canary: fmt.Sprintf("%s.%s", payload.Name, payload.Namespace),
+						logger: logger,
+					},
+				}
+
+				ctx, cancel := context.WithTimeout(context.Background(), taskRunner.Timeout())
+				defer cancel()
+
+				result, err := helm.Run(ctx)
+				if !result.ok {
+					w.WriteHeader(http.StatusInternalServerError)
+					w.Write([]byte(err.Error()))
+					return
+				}
+
+				w.WriteHeader(http.StatusOK)
+				if rtnCmdOutput {
+					w.Write(result.out)
+				}
+				return
+			}
+
+			// run concord job (blocking task)
+			if typ == TaskTypeConcord {
+				concord, err := NewConcordTask(payload.Metadata, fmt.Sprintf("%s.%s", payload.Name, payload.Namespace), logger)
+
+				if err != nil {
+					logger.With("canary", payload.Name).Errorf("concord task init error: %s", err)
+					w.WriteHeader(http.StatusInternalServerError)
+					w.Write([]byte(err.Error()))
+					return
+				}
+
+				ctx, cancel := context.WithTimeout(context.Background(), taskRunner.Timeout())
+				defer cancel()
+
+				result, err := concord.Run(ctx)
+				if !result.ok {
+					if err != nil {
+						logger.With("canary", payload.Name).Errorf("concord task error: %s", err)
+					}
+					w.WriteHeader(http.StatusInternalServerError)
+					w.Write([]byte(err.Error()))
+					return
+				}
+
+				w.WriteHeader(http.StatusOK)
+				if rtnCmdOutput {
+					w.Write(result.out)
+				}
+				return
+			}
+
+			taskFactory, ok := GetTaskFactory(typ)
+			if !ok {
+				w.WriteHeader(http.StatusBadRequest)
+				w.Write([]byte(fmt.Sprintf("unknown task type %s", typ)))
+				return
+			}
+			canary := fmt.Sprintf("%s.%s", payload.Name, payload.Namespace)
+			task, err := taskFactory(metadata, canary, logger)
+			if err != nil {
+				w.WriteHeader(http.StatusBadRequest)
+				w.Write([]byte(err.Error()))
+				return
+			}
+			taskRunner.Add(task)
+		} else {
+			w.WriteHeader(http.StatusBadRequest)
+			w.Write([]byte("metadata not found in payload"))
+			return
+		}
+
+		w.WriteHeader(http.StatusAccepted)
+	}
+}

pkg/loadtester/server_fixture_test.go

@@ -0,0 +1,46 @@
+package loadtester
+
+import (
+	"net/http/httptest"
+	"time"
+
+	"github.com/weaveworks/flagger/pkg/logger"
+	"go.uber.org/zap"
+)
+
+type serverFixture struct {
+	taskRunner *MockTaskRunner
+	resp       *httptest.ResponseRecorder
+	logger     *zap.SugaredLogger
+}
+
+func newServerFixture() serverFixture {
+	taskRunner := &MockTaskRunner{}
+	resp := httptest.NewRecorder()
+	logger, _ := logger.NewLogger("info")
+
+	return serverFixture{
+		taskRunner: taskRunner,
+		resp:       resp,
+		logger:     logger,
+	}
+}
+
+type MockTaskRunner struct {
+}
+
+func (m *MockTaskRunner) Add(task Task) {
+
+}
+
+func (m *MockTaskRunner) GetTotalExecs() uint64 {
+	return 0
+}
+
+func (m *MockTaskRunner) Start(interval time.Duration, stopCh <-chan struct{}) {
+
+}
+
+func (m *MockTaskRunner) Timeout() time.Duration {
+	return time.Hour
+}

pkg/loadtester/server_test.go

@@ -0,0 +1,75 @@
+package loadtester
+
+import (
+	"bytes"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	flaggerv1 "github.com/weaveworks/flagger/pkg/apis/flagger/v1beta1"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestServer_HandleHealthz(t *testing.T) {
+	req, _ := http.NewRequest("GET", "/heathz", nil)
+	resp := httptest.NewRecorder()
+	HandleHealthz(resp, req)
+
+	assert.Equal(t, http.StatusOK, resp.Code)
+	assert.Equal(t, "OK", resp.Body.String())
+}
+
+func TestServer_HandleNewBashTaskCmdExitZero(t *testing.T) {
+	mocks := newServerFixture()
+	resp := mocks.resp
+	req := newJsonRequest("POST", "/", &flaggerv1.CanaryWebhookPayload{
+		Metadata: map[string]string{
+			"type": TaskTypeBash,
+			"cmd":  "echo some-output-not-to-be-returned",
+		},
+	})
+	HandleNewTask(mocks.logger, mocks.taskRunner)(resp, req)
+
+	assert.Equal(t, http.StatusOK, resp.Code)
+	assert.Empty(t, resp.Body.String())
+}
+
+func TestServer_HandleNewBashTaskCmdExitZeroReturnCmdOutput(t *testing.T) {
+	mocks := newServerFixture()
+	resp := mocks.resp
+	req := newJsonRequest("POST", "/", &flaggerv1.CanaryWebhookPayload{
+		Metadata: map[string]string{
+			"type":            TaskTypeBash,
+			"cmd":             "echo some-output-to-be-returned",
+			"returnCmdOutput": "true",
+		},
+	})
+	HandleNewTask(mocks.logger, mocks.taskRunner)(resp, req)
+
+	assert.Equal(t, http.StatusOK, resp.Code)
+	assert.Equal(t, "some-output-to-be-returned\n", resp.Body.String())
+}
+
+func TestServer_HandleNewBashTaskCmdExitNonZero(t *testing.T) {
+	mocks := newServerFixture()
+	resp := mocks.resp
+	req := newJsonRequest("POST", "/", &flaggerv1.CanaryWebhookPayload{
+		Metadata: map[string]string{
+			"type": TaskTypeBash,
+			"cmd":  "false",
+		},
+	})
+
+	HandleNewTask(mocks.logger, mocks.taskRunner)(resp, req)
+
+	assert.Equal(t, http.StatusInternalServerError, resp.Code)
+	assert.Equal(t, "command false failed: : exit status 1", resp.Body.String())
+}
+
+func newJsonRequest(method string, url string, v interface{}) *http.Request {
+	payload, _ := json.Marshal(v)
+	req, _ := http.NewRequest(method, url, bytes.NewReader(payload))
+	return req
+}

pkg/loadtester/task.go

@@ -12,7 +12,7 @@ import (
 // Modeling a loadtester task
 type Task interface {
 	Hash() string
-	Run(ctx context.Context) bool
+	Run(ctx context.Context) *TaskRunResult
 	String() string
 	Canary() string
 }
@@ -40,3 +40,8 @@ func GetTaskFactory(typ string) (TaskFactory, bool) {
 	factory, ok := taskFactories.Load(typ)
 	return factory.(TaskFactory), ok
 }
+
+type TaskRunResult struct {
+	ok  bool
+	out []byte
+}

pkg/loadtester/task_ngrinder.go

@@ -86,17 +86,17 @@ func (task *NGrinderTask) StopEndpoint() *url.URL {
 }
 
 // initiate a clone_and_start request and get new test id from response
-func (task *NGrinderTask) Run(ctx context.Context) bool {
+func (task *NGrinderTask) Run(ctx context.Context) *TaskRunResult {
 	url := task.CloneAndStartEndpoint().String()
 	result, err := task.request("POST", url, ctx)
 	if err != nil {
 		task.logger.With("canary", task.canary).
 			Errorf("failed to clone and start ngrinder test %s: %s", url, err.Error())
-		return false
+		return &TaskRunResult{false, nil}
 	}
 	id := result["id"]
 	task.testId = int(id.(float64))
-	return task.PollStatus(ctx)
+	return &TaskRunResult{task.PollStatus(ctx), nil}
 }
 
 func (task *NGrinderTask) String() string {

pkg/loadtester/task_shell.go

@@ -33,7 +33,7 @@ func (task *CmdTask) Hash() string {
 	return hash(task.canary + task.command)
 }
 
-func (task *CmdTask) Run(ctx context.Context) bool {
+func (task *CmdTask) Run(ctx context.Context) *TaskRunResult {
 	cmd := exec.CommandContext(ctx, "sh", "-c", task.command)
 	out, err := cmd.CombinedOutput()
 
@@ -45,7 +45,7 @@ func (task *CmdTask) Run(ctx context.Context) bool {
 		}
 		task.logger.With("canary", task.canary).Infof("command finished %s", task.command)
 	}
-	return err == nil
+	return &TaskRunResult{err == nil, out}
 }
 
 func (task *CmdTask) String() string {