From eed870ff0501f1a8d605c18a05343cdab7ec8b84 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ivan=20Miri=C4=87?= <ivan@loadimpact.com>
Date: Wed, 17 Feb 2021 12:39:03 +0100
Subject: [PATCH] Add getTestStats to k6/execution

Part of #1320
---
 core/local/local.go                  |  1 +
 js/modules/k6/execution/execution.go | 20 ++++++++++++++++++++
 js/runner_test.go                    | 21 +++++++++++++++++++++
 lib/context.go                       | 17 +++++++++++++++++
 4 files changed, 59 insertions(+)

diff --git a/core/local/local.go b/core/local/local.go
index ed1d9d66232..09b4939aba0 100644
--- a/core/local/local.go
+++ b/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...
 
diff --git a/js/modules/k6/execution/execution.go b/js/modules/k6/execution/execution.go
index a68ebbbf2b4..a55e3ca468a 100644
--- a/js/modules/k6/execution/execution.go
+++ b/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
+}
diff --git a/js/runner_test.go b/js/runner_test.go
index 82e6608cd41..0448d717af5 100644
--- a/js/runner_test.go
+++ b/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)
 		})
diff --git a/lib/context.go b/lib/context.go
index 97681f1cb75..64df24f3942 100644
--- a/lib/context.go
+++ b/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)