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

@@ -72,3 +72,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(),
+		"iterations_completed":   es.GetFullIterationCount(),
+		"iterations_interrupted": es.GetPartialIterationCount(),
+		"vus_active":             es.GetCurrentlyActiveVUsCount(),
+		"vus_max":                es.GetInitializedVUsCount(),
+	}
+
+	return out, nil
+}

js/runner_test.go

@@ -2038,6 +2038,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.vus_active !== 1) throw new Error('unexpected vus_active: '+ts.vus_active);
+			if (ts.vus_max !== 0) throw new Error('unexpected vus_max: '+ts.vus_max);
+			if (ts.iterations_completed !== 0) throw new Error('unexpected iterations_completed: '+ts.iterations_completed);
+			if (ts.iterations_interrupted !== 0) throw new Error('unexpected iterations_interrupted: '+ts.iterations_interrupted);
+		}`, ""},
+		{"test_err", `
+		var exec = require('k6/execution');
+		exec.getTestStats();
+		`, "test information can only be returned from an exported function"},
 	}
 
 	for _, tc := range testCases {
@@ -2055,8 +2069,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())
+
 			vu := initVU.Activate(&lib.VUActivationParams{
 				RunContext: ctx,
 				Exec:       "default",
@@ -2070,6 +2089,8 @@ func TestExecutionStats(t *testing.T) {
 				},
 			})
 
+			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)