Add methods for setup and teardown control in the API

This would hopefully be enough for fully implementing #539 in the backend

api/v1/routes.go

@@ -38,5 +38,11 @@ func NewHandler() http.Handler {
 	router.GET("/v1/groups", HandleGetGroups)
 	router.GET("/v1/groups/:id", HandleGetGroup)
 
+	router.POST("/v1/setup", HandleRunSetup)
+	router.PUT("/v1/setup", HandleSetSetupData)
+	router.GET("/v1/setup", HandleGetSetupData)
+
+	router.POST("/v1/teardown", HandleRunTeardown)
+
 	return router
 }

api/v1/setup_teardown_routes.go

@@ -0,0 +1,104 @@
+/*
+ *
+ * k6 - a next-generation load testing tool
+ * Copyright (C) 2018 Load Impact
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as
+ * published by the Free Software Foundation, either version 3 of the
+ * License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ */
+
+package v1
+
+import (
+	"encoding/json"
+	"io/ioutil"
+	"net/http"
+
+	"github.com/julienschmidt/httprouter"
+	"github.com/loadimpact/k6/api/common"
+	"github.com/manyminds/api2go/jsonapi"
+)
+
+// SetupData is just a simple wrapper to satisfy jsonapi
+type SetupData struct {
+	Data interface{} `json:"data" yaml:"data"`
+}
+
+// GetName is a dummy method so we can satisfy the jsonapi.EntityNamer interface
+func (sd SetupData) GetName() string {
+	return "setupData"
+}
+
+// GetID is a dummy method so we can satisfy the jsonapi.MarshalIdentifier interface
+func (sd SetupData) GetID() string {
+	return "default"
+}
+
+func handleSetupDataOutput(rw http.ResponseWriter, setupData interface{}) {
+	rw.Header().Set("Content-Type", "application/json")
+
+	data, err := jsonapi.Marshal(SetupData{setupData})
+	if err != nil {
+		apiError(rw, "Encoding error", err.Error(), http.StatusInternalServerError)
+		return
+	}
+	_, _ = rw.Write(data)
+}
+
+// HandleGetSetupData just returns the current JSON-encoded setup data
+func HandleGetSetupData(rw http.ResponseWriter, r *http.Request, p httprouter.Params) {
+	runner := common.GetEngine(r.Context()).Executor.GetRunner()
+	handleSetupDataOutput(rw, runner.GetSetupData())
+}
+
+// HandleSetSetupData just parses the JSON request body and sets the result as setup data for the runner
+func HandleSetSetupData(rw http.ResponseWriter, r *http.Request, p httprouter.Params) {
+	body, err := ioutil.ReadAll(r.Body)
+	if err != nil {
+		apiError(rw, "Error reading request body", err.Error(), http.StatusBadRequest)
+		return
+	}
+
+	var setupData interface{}
+	if err := json.Unmarshal(body, &setupData); err != nil {
+		apiError(rw, "Error parsing request body", err.Error(), http.StatusBadRequest)
+		return
+	}
+
+	runner := common.GetEngine(r.Context()).Executor.GetRunner()
+	runner.SetSetupData(setupData)
+
+	handleSetupDataOutput(rw, runner.GetSetupData())
+}
+
+// HandleRunSetup executes the runner's Setup() method and returns the result
+func HandleRunSetup(rw http.ResponseWriter, r *http.Request, p httprouter.Params) {
+	runner := common.GetEngine(r.Context()).Executor.GetRunner()
+
+	if err := runner.Setup(r.Context()); err != nil {
+		apiError(rw, "Error executing setup", err.Error(), http.StatusInternalServerError)
+		return
+	}
+
+	handleSetupDataOutput(rw, runner.GetSetupData())
+}
+
+// HandleRunTeardown executes the runner's Teardown() method
+func HandleRunTeardown(rw http.ResponseWriter, r *http.Request, p httprouter.Params) {
+	runner := common.GetEngine(r.Context()).Executor.GetRunner()
+
+	if err := runner.Teardown(r.Context()); err != nil {
+		apiError(rw, "Error executing teardown", err.Error(), http.StatusInternalServerError)
+	}
+}

api/v1/setup_teardown_routes_test.go

@@ -0,0 +1,119 @@
+/*
+ *
+ * k6 - a next-generation load testing tool
+ * Copyright (C) 2018 Load Impact
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Affero General Public License as
+ * published by the Free Software Foundation, either version 3 of the
+ * License, or (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ */
+
+package v1
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/loadimpact/k6/core"
+	"github.com/loadimpact/k6/core/local"
+	"github.com/loadimpact/k6/js"
+	"github.com/loadimpact/k6/lib"
+	"github.com/loadimpact/k6/lib/types"
+	"github.com/manyminds/api2go/jsonapi"
+	"github.com/spf13/afero"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	null "gopkg.in/guregu/null.v3"
+)
+
+func TestSetupData(t *testing.T) {
+	t.Parallel()
+	runner, err := js.New(
+		&lib.SourceData{Filename: "/script.js", Data: []byte(`
+			export function setup() {
+				return {"v": 1};
+			}
+
+			export default function(data) {
+				if (!data || data.v != 2) {
+					throw new Error("incorrect data: " + JSON.stringify(data));
+				}
+			};
+
+			export function teardown(data) {
+				if (!data || data.v != 2) {
+					throw new Error("incorrect teardown data: " + JSON.stringify(data));
+				}
+			}
+
+		`)},
+		afero.NewMemMapFs(),
+		lib.RuntimeOptions{},
+	)
+	require.NoError(t, err)
+	runner.SetOptions(lib.Options{
+		Paused:          null.BoolFrom(true),
+		VUs:             null.IntFrom(2),
+		VUsMax:          null.IntFrom(2),
+		Iterations:      null.IntFrom(3),
+		SetupTimeout:    types.NullDurationFrom(1 * time.Second),
+		TeardownTimeout: types.NullDurationFrom(1 * time.Second),
+	})
+	executor := local.New(runner)
+	executor.SetRunSetup(false)
+	engine, err := core.NewEngine(executor, runner.GetOptions())
+	require.NoError(t, err)
+
+	handler := NewHandler()
+
+	checkSetup := func(method, body, expResult string) {
+		rw := httptest.NewRecorder()
+		handler.ServeHTTP(rw, newRequestWithEngine(engine, method, "/v1/setup", bytes.NewBufferString(body)))
+		res := rw.Result()
+		assert.Equal(t, http.StatusOK, res.StatusCode)
+
+		var doc jsonapi.Document
+		assert.NoError(t, json.Unmarshal(rw.Body.Bytes(), &doc))
+		if !assert.NotNil(t, doc.Data.DataObject) {
+			return
+		}
+		assert.Equal(t, "setupData", doc.Data.DataObject.Type)
+		assert.JSONEq(t, expResult, string(doc.Data.DataObject.Attributes))
+	}
+
+	checkSetup("GET", "", `{"data": null}`)
+	checkSetup("POST", "", `{"data": {"v":1}}`)
+	checkSetup("GET", "", `{"data": {"v":1}}`)
+	checkSetup("PUT", `{"v":2, "test":"mest"}`, `{"data": {"v":2, "test":"mest"}}`)
+	checkSetup("GET", "", `{"data": {"v":2, "test":"mest"}}`)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	errC := make(chan error)
+	go func() { errC <- engine.Run(ctx) }()
+
+	engine.Executor.SetPaused(false)
+
+	select {
+	case <-time.After(10 * time.Second):
+		cancel()
+		t.Fatal("Test timed out")
+	case err := <-errC:
+		cancel()
+		require.NoError(t, err)
+	}
+}

core/local/local_test.go

@@ -54,9 +54,9 @@ func TestExecutorSetupTeardownRun(t *testing.T) {
 		setupC := make(chan struct{})
 		teardownC := make(chan struct{})
 		e := New(&lib.MiniRunner{
-			SetupFn: func(ctx context.Context) error {
+			SetupFn: func(ctx context.Context) (interface{}, error) {
 				close(setupC)
-				return nil
+				return nil, nil
 			},
 			TeardownFn: func(ctx context.Context) error {
 				close(teardownC)
@@ -74,8 +74,8 @@ func TestExecutorSetupTeardownRun(t *testing.T) {
 	})
 	t.Run("Setup Error", func(t *testing.T) {
 		e := New(&lib.MiniRunner{
-			SetupFn: func(ctx context.Context) error {
-				return errors.New("setup error")
+			SetupFn: func(ctx context.Context) (interface{}, error) {
+				return nil, errors.New("setup error")
 			},
 			TeardownFn: func(ctx context.Context) error {
 				return errors.New("teardown error")
@@ -85,8 +85,8 @@ func TestExecutorSetupTeardownRun(t *testing.T) {
 
 		t.Run("Don't Run Setup", func(t *testing.T) {
 			e := New(&lib.MiniRunner{
-				SetupFn: func(ctx context.Context) error {
-					return errors.New("setup error")
+				SetupFn: func(ctx context.Context) (interface{}, error) {
+					return nil, errors.New("setup error")
 				},
 				TeardownFn: func(ctx context.Context) error {
 					return errors.New("teardown error")
@@ -101,8 +101,8 @@ func TestExecutorSetupTeardownRun(t *testing.T) {
 	})
 	t.Run("Teardown Error", func(t *testing.T) {
 		e := New(&lib.MiniRunner{
-			SetupFn: func(ctx context.Context) error {
-				return nil
+			SetupFn: func(ctx context.Context) (interface{}, error) {
+				return nil, nil
 			},
 			TeardownFn: func(ctx context.Context) error {
 				return errors.New("teardown error")
@@ -115,8 +115,8 @@ func TestExecutorSetupTeardownRun(t *testing.T) {
 
 		t.Run("Don't Run Teardown", func(t *testing.T) {
 			e := New(&lib.MiniRunner{
-				SetupFn: func(ctx context.Context) error {
-					return nil
+				SetupFn: func(ctx context.Context) (interface{}, error) {
+					return nil, nil
 				},
 				TeardownFn: func(ctx context.Context) error {
 					return errors.New("teardown error")

js/runner.go

@@ -46,6 +46,9 @@ import (
 
 var errInterrupt = errors.New("context cancelled")
 
+// Ensure Runner implements the lib.Runner interface
+var _ lib.Runner = &Runner{}
+
 type Runner struct {
 	Bundle       *Bundle
 	Logger       *log.Logger
@@ -197,6 +200,16 @@ func (r *Runner) Setup(ctx context.Context) error {
 	return json.Unmarshal(data, &r.setupData)
 }
 
+// GetSetupData returns the setup data if Setup() was specified and executed, nil otherwise
+func (r *Runner) GetSetupData() interface{} {
+	return r.setupData
+}
+
+// SetSetupData saves the externally supplied setup data in the runner, so it can be used in VUs
+func (r *Runner) SetSetupData(data interface{}) {
+	r.setupData = data
+}
+
 func (r *Runner) Teardown(ctx context.Context) error {
 	_, err := r.runPart(ctx, "teardown", r.setupData)
 	return err

lib/runner.go

@@ -48,6 +48,12 @@ type Runner interface {
 	// Runs pre-test setup, if applicable.
 	Setup(ctx context.Context) error
 
+	// Returns the setup data if setup() is specified and run, nil otherwise
+	GetSetupData() interface{}
+
+	// Saves the externally supplied setup data in the runner
+	SetSetupData(interface{})
+
 	// Runs post-test teardown, if applicable.
 	Teardown(ctx context.Context) error
 
@@ -75,9 +81,11 @@ type VU interface {
 // MiniRunner wraps a function in a runner whose VUs will simply call that function.
 type MiniRunner struct {
 	Fn         func(ctx context.Context) ([]stats.SampleContainer, error)
-	SetupFn    func(ctx context.Context) error
+	SetupFn    func(ctx context.Context) (interface{}, error)
 	TeardownFn func(ctx context.Context) error
 
+	setupData interface{}
+
 	Group   *Group
 	Options Options
 }
@@ -96,11 +104,23 @@ func (r MiniRunner) NewVU() (VU, error) {
 
 func (r MiniRunner) Setup(ctx context.Context) error {
 	if fn := r.SetupFn; fn != nil {
-		return fn(ctx)
+		data, err := fn(ctx)
+		r.setupData = data
+		return err
 	}
 	return nil
 }
 
+// GetSetupData returns the setup data if Setup() was executed, nil otherwise
+func (r MiniRunner) GetSetupData() interface{} {
+	return r.setupData
+}
+
+// SetSetupData saves the externally supplied setup data in the runner
+func (r MiniRunner) SetSetupData(data interface{}) {
+	r.setupData = data
+}
+
 func (r MiniRunner) Teardown(ctx context.Context) error {
 	if fn := r.TeardownFn; fn != nil {
 		return fn(ctx)