js/k6: ModuleV2 migration

myjeffxie · Nov 22, 2021 · e9ecfe7 · e9ecfe7
1 parent 795d937
commit e9ecfe7
Show file tree

Hide file tree

Showing 2 changed files with 120 additions and 34 deletions.
diff --git a/js/modules/k6/k6.go b/js/modules/k6/k6.go
@@ -22,7 +22,6 @@
 package k6
 
 import (
-	"context"
 	"errors"
 	"math/rand"
 	"sync/atomic"
@@ -31,22 +30,56 @@ import (
 	"github.com/dop251/goja"
 
 	"go.k6.io/k6/js/common"
-	"go.k6.io/k6/lib"
+	"go.k6.io/k6/js/modules"
 	"go.k6.io/k6/stats"
 )
 
-// K6 is just the module struct.
-type K6 struct{}
+var (
+	// ErrGroupInInitContext is returned when group() are using in the init context.
+	ErrGroupInInitContext = common.NewInitContextError("Using group() in the init context is not supported")
 
-// ErrGroupInInitContext is returned when group() are using in the init context.
-var ErrGroupInInitContext = common.NewInitContextError("Using group() in the init context is not supported")
+	// ErrCheckInInitContext is returned when check() are using in the init context.
+	ErrCheckInInitContext = common.NewInitContextError("Using check() in the init context is not supported")
+)
+
+type (
+	// RootModule is the global module instance that will create module
+	// instances for each VU.
+	RootModule struct{}
+
+	// K6 represents an instance of the k6 module.
+	K6 struct {
+		vu modules.VU
+	}
+)
+
+var (
+	_ modules.Module   = &RootModule{}
+	_ modules.Instance = &K6{}
+)
+
+// New returns a pointer to a new RootModule instance.
+func New() *RootModule {
+	return &RootModule{}
+}
 
-// ErrCheckInInitContext is returned when check() are using in the init context.
-var ErrCheckInInitContext = common.NewInitContextError("Using check() in the init context is not supported")
+// NewModuleInstance implements the modules.Module interface to return
+// a new instance for each VU.
+func (*RootModule) NewModuleInstance(vu modules.VU) modules.Instance {
+	return &K6{vu: vu}
+}
 
-// New returns a new module Struct.
-func New() *K6 {
-	return &K6{}
+// Exports returns the exports of the k6 module.
+func (mi *K6) Exports() modules.Exports {
+	return modules.Exports{
+		Named: map[string]interface{}{
+			"check":      mi.Check,
+			"fail":       mi.Fail,
+			"group":      mi.Group,
+			"randomSeed": mi.RandomSeed,
+			"sleep":      mi.Sleep,
+		},
+	}
 }
 
 // Fail is a fancy way of saying `throw "something"`.
@@ -55,7 +88,8 @@ func (*K6) Fail(msg string) (goja.Value, error) {
 }
 
 // Sleep waits the provided seconds before continuing the execution.
-func (*K6) Sleep(ctx context.Context, secs float64) {
+func (mi *K6) Sleep(secs float64) {
+	ctx := mi.vu.Context()
 	timer := time.NewTimer(time.Duration(secs * float64(time.Second)))
 	select {
 	case <-timer.C:
@@ -65,16 +99,14 @@ func (*K6) Sleep(ctx context.Context, secs float64) {
 }
 
 // RandomSeed sets the seed to the random generator used for this VU.
-func (*K6) RandomSeed(ctx context.Context, seed int64) {
+func (mi *K6) RandomSeed(seed int64) {
 	randSource := rand.New(rand.NewSource(seed)).Float64 //nolint:gosec
-
-	rt := common.GetRuntime(ctx)
-	rt.SetRandSource(randSource)
+	mi.vu.Runtime().SetRandSource(randSource)
 }
 
 // Group wraps a function call and executes it within the provided group name.
-func (*K6) Group(ctx context.Context, name string, fn goja.Callable) (goja.Value, error) {
-	state := lib.GetState(ctx)
+func (mi *K6) Group(name string, fn goja.Callable) (goja.Value, error) {
+	state := mi.vu.State()
 	if state == nil {
 		return nil, ErrGroupInInitContext
 	}
@@ -108,6 +140,7 @@ func (*K6) Group(ctx context.Context, name string, fn goja.Callable) (goja.Value
 
 	tags := state.CloneTags()
 
+	ctx := mi.vu.Context()
 	stats.PushIfNotDone(ctx, state.Samples, stats.Sample{
 		Time:   t,
 		Metric: state.BuiltinMetrics.GroupDuration,
@@ -120,12 +153,13 @@ func (*K6) Group(ctx context.Context, name string, fn goja.Callable) (goja.Value
 
 // Check will emit check metrics for the provided checks.
 //nolint:cyclop
-func (*K6) Check(ctx context.Context, arg0, checks goja.Value, extras ...goja.Value) (bool, error) {
-	state := lib.GetState(ctx)
+func (mi *K6) Check(arg0, checks goja.Value, extras ...goja.Value) (bool, error) {
+	state := mi.vu.State()
 	if state == nil {
 		return false, ErrCheckInInitContext
 	}
-	rt := common.GetRuntime(ctx)
+	ctx := mi.vu.Context()
+	rt := mi.vu.Runtime()
 	t := time.Now()
 
 	// Prepare the metric tags

diff --git a/js/modules/k6/k6_test.go b/js/modules/k6/k6_test.go
@@ -32,6 +32,7 @@ import (
 	"github.com/stretchr/testify/require"
 
 	"go.k6.io/k6/js/common"
+	"go.k6.io/k6/js/modulestest"
 	"go.k6.io/k6/lib"
 	"go.k6.io/k6/lib/metrics"
 	"go.k6.io/k6/stats"
@@ -40,7 +41,17 @@ import (
 func TestFail(t *testing.T) {
 	t.Parallel()
 	rt := goja.New()
-	require.NoError(t, rt.Set("k6", common.Bind(rt, New(), nil)))
+	m, ok := New().NewModuleInstance(
+		&modulestest.VU{
+			RuntimeField: rt,
+			InitEnvField: &common.InitEnvironment{},
+			CtxField:     common.WithRuntime(context.Background(), rt),
+			StateField:   nil,
+		},
+	).(*K6)
+	require.True(t, ok)
+	require.NoError(t, rt.Set("k6", m.Exports().Named))
+
 	_, err := rt.RunString(`k6.fail("blah")`)
 	assert.Contains(t, err.Error(), "blah")
 }
@@ -58,8 +69,17 @@ func TestSleep(t *testing.T) {
 		t.Run(name, func(t *testing.T) {
 			t.Parallel()
 			rt := goja.New()
-			ctx := context.Background()
-			require.NoError(t, rt.Set("k6", common.Bind(rt, New(), &ctx)))
+			m, ok := New().NewModuleInstance(
+				&modulestest.VU{
+					RuntimeField: rt,
+					InitEnvField: &common.InitEnvironment{},
+					CtxField:     common.WithRuntime(context.Background(), rt),
+					StateField:   nil,
+				},
+			).(*K6)
+			require.True(t, ok)
+			require.NoError(t, rt.Set("k6", m.Exports().Named))
+
 			startTime := time.Now()
 			_, err := rt.RunString(`k6.sleep(1)`)
 			endTime := time.Now()
@@ -72,7 +92,17 @@ func TestSleep(t *testing.T) {
 		t.Parallel()
 		rt := goja.New()
 		ctx, cancel := context.WithCancel(context.Background())
-		require.NoError(t, rt.Set("k6", common.Bind(rt, New(), &ctx)))
+		m, ok := New().NewModuleInstance(
+			&modulestest.VU{
+				RuntimeField: rt,
+				InitEnvField: &common.InitEnvironment{},
+				CtxField:     common.WithRuntime(ctx, rt),
+				StateField:   nil,
+			},
+		).(*K6)
+		require.True(t, ok)
+		require.NoError(t, rt.Set("k6", m.Exports().Named))
+
 		dch := make(chan time.Duration)
 		go func() {
 			startTime := time.Now()
@@ -96,10 +126,16 @@ func TestRandSeed(t *testing.T) {
 	t.Parallel()
 	rt := goja.New()
 
-	ctx := context.Background()
-	ctx = common.WithRuntime(ctx, rt)
-
-	require.NoError(t, rt.Set("k6", common.Bind(rt, New(), &ctx)))
+	m, ok := New().NewModuleInstance(
+		&modulestest.VU{
+			RuntimeField: rt,
+			InitEnvField: &common.InitEnvironment{},
+			CtxField:     common.WithRuntime(context.Background(), rt),
+			StateField:   nil,
+		},
+	).(*K6)
+	require.True(t, ok)
+	require.NoError(t, rt.Set("k6", m.Exports().Named))
 
 	rand := 0.8487305991992138
 	_, err := rt.RunString(fmt.Sprintf(`
@@ -131,11 +167,18 @@ func TestGroup(t *testing.T) {
 				SystemTags: stats.NewSystemTagSet(stats.TagGroup),
 			},
 		}
-		ctx := context.Background()
-		ctx = lib.WithState(ctx, state)
-		ctx = common.WithRuntime(ctx, rt)
+		ctx := lib.WithState(context.Background(), state)
 		state.BuiltinMetrics = metrics.RegisterBuiltinMetrics(metrics.NewRegistry())
-		require.NoError(t, rt.Set("k6", common.Bind(rt, New(), &ctx)))
+		m, ok := New().NewModuleInstance(
+			&modulestest.VU{
+				RuntimeField: rt,
+				InitEnvField: &common.InitEnvironment{},
+				CtxField:     common.WithRuntime(ctx, rt),
+				StateField:   state,
+			},
+		).(*K6)
+		require.True(t, ok)
+		require.NoError(t, rt.Set("k6", m.Exports().Named))
 		return rt, state, root
 	}
 
@@ -191,7 +234,16 @@ func checkTestRuntime(t testing.TB, ctxs ...*context.Context) (
 	ctx = common.WithRuntime(ctx, rt)
 	ctx = lib.WithState(ctx, state)
 	state.BuiltinMetrics = metrics.RegisterBuiltinMetrics(metrics.NewRegistry())
-	require.NoError(t, rt.Set("k6", common.Bind(rt, New(), &ctx)))
+	m, ok := New().NewModuleInstance(
+		&modulestest.VU{
+			RuntimeField: rt,
+			InitEnvField: &common.InitEnvironment{},
+			CtxField:     ctx,
+			StateField:   state,
+		},
+	).(*K6)
+	require.True(t, ok)
+	require.NoError(t, rt.Set("k6", m.Exports().Named))
 	if len(ctxs) == 1 { // hacks
 		*ctxs[0] = ctx
 	}