Add getTestStats to k6/execution

Part of #1320

core/local/local.go

@@ -361,6 +361,7 @@ func (e *ExecutionScheduler) Run(globalCtx, runCtx context.Context, engineOut ch
 
 	runResults := make(chan error, executorsCount) // nil values are successful runs
 
+	runCtx = lib.WithExecutionState(runCtx, e.state)
 	runSubCtx, cancel := context.WithCancel(runCtx)
 	defer cancel() // just in case, and to shut up go vet...
 

js/modules/k6/execution/execution.go

@@ -67,3 +67,23 @@ func (e *Execution) GetScenarioStats(ctx context.Context) (map[string]interface{
 
 	return out, nil
 }
+
+// GetTestStats returns global test information.
+func (e *Execution) GetTestStats(ctx context.Context) (map[string]interface{}, error) {
+	es := lib.GetExecutionState(ctx)
+	if es == nil {
+		return nil, errors.New("test information can only be returned from an exported function")
+	}
+
+	out := map[string]interface{}{
+		// XXX: For consistency, should this be startTime instead, or startTime
+		// in ScenarioStats be converted to duration?
+		"duration":              es.GetCurrentTestRunDuration().String(),
+		"iterationsCompleted":   es.GetFullIterationCount(),
+		"iterationsInterrupted": es.GetPartialIterationCount(),
+		"vusActive":             es.GetCurrentlyActiveVUsCount(),
+		"vusMax":                es.GetInitializedVUsCount(),
+	}
+
+	return out, nil
+}

js/runner_test.go

@@ -2058,6 +2058,20 @@ func TestExecutionStats(t *testing.T) {
 		var exec = require('k6/execution');
 		exec.getScenarioStats();
 		`, "scenario information can only be returned from an exported function"},
+		{"test_ok", `
+		var exec = require('k6/execution');
+
+		exports.default = function() {
+			var ts = exec.getTestStats();
+			if (ts.vusActive !== 1) throw new Error('unexpected vusActive: '+ts.vusActive);
+			if (ts.vusMax !== 0) throw new Error('unexpected vusMax: '+ts.vusMax);
+			if (ts.iterationsCompleted !== 0) throw new Error('unexpected iterationsCompleted: '+ts.iterationsCompleted);
+			if (ts.iterationsInterrupted !== 0) throw new Error('unexpected iterationsInterrupted: '+ts.iterationsInterrupted);
+		}`, ""},
+		{"test_err", `
+		var exec = require('k6/execution');
+		exec.getTestStats();
+		`, "test information can only be returned from an exported function"},
 	}
 
 	for _, tc := range testCases {
@@ -2075,9 +2089,13 @@ func TestExecutionStats(t *testing.T) {
 			initVU, err := r.newVU(1, samples)
 			require.NoError(t, err)
 
+			execScheduler, err := local.NewExecutionScheduler(r, testutils.NewLogger(t))
+			require.NoError(t, err)
+
 			ctx, cancel := context.WithCancel(context.Background())
 			defer cancel()
 
+			ctx = lib.WithExecutionState(ctx, execScheduler.GetState())
 			ctx = lib.WithScenarioState(ctx, &lib.ScenarioState{
 				Name:      "default",
 				Executor:  "test-exec",
@@ -2091,6 +2109,8 @@ func TestExecutionStats(t *testing.T) {
 				Exec:       "default",
 			})
 
+			execState := execScheduler.GetState()
+			execState.ModCurrentlyActiveVUsCount(+1)
 			err = vu.RunOnce()
 			assert.NoError(t, err)
 		})

lib/context.go

@@ -26,13 +26,16 @@ type ctxKey int
 
 const (
 	ctxKeyState ctxKey = iota
+	ctxKeyExecState
 	ctxKeyScenario
 )
 
+// WithState embeds a State in ctx.
 func WithState(ctx context.Context, state *State) context.Context {
 	return context.WithValue(ctx, ctxKeyState, state)
 }
 
+// GetState returns a State from ctx.
 func GetState(ctx context.Context) *State {
 	v := ctx.Value(ctxKeyState)
 	if v == nil {
@@ -41,6 +44,20 @@ func GetState(ctx context.Context) *State {
 	return v.(*State)
 }
 
+// WithExecutionState embeds an ExecutionState in ctx.
+func WithExecutionState(ctx context.Context, s *ExecutionState) context.Context {
+	return context.WithValue(ctx, ctxKeyExecState, s)
+}
+
+// GetExecutionState returns an ExecutionState from ctx.
+func GetExecutionState(ctx context.Context) *ExecutionState {
+	v := ctx.Value(ctxKeyExecState)
+	if v == nil {
+		return nil
+	}
+	return v.(*ExecutionState)
+}
+
 // WithScenarioState embeds a ScenarioState in ctx.
 func WithScenarioState(ctx context.Context, s *ScenarioState) context.Context {
 	return context.WithValue(ctx, ctxKeyScenario, s)