Add testmodules.Runtime and use it in k6/ws and k6/grpc tests

fixes #2598

js/modules/k6/grpc/client_test.go

@@ -61,36 +61,16 @@ func TestClient(t *testing.T) {
 	t.Parallel()
 
 	type testState struct {
-		rt      *goja.Runtime
-		vuState *lib.State
-		env     *common.InitEnvironment
+		*modulestest.Runtime
 		httpBin *httpmultibin.HTTPMultiBin
 		samples chan metrics.SampleContainer
 	}
 	setup := func(t *testing.T) testState {
 		t.Helper()
 
-		root, err := lib.NewGroup("", nil)
-		require.NoError(t, err)
 		tb := httpmultibin.NewHTTPMultiBin(t)
 		samples := make(chan metrics.SampleContainer, 1000)
-		state := &lib.State{
-			Group:     root,
-			Dialer:    tb.Dialer,
-			TLSConfig: tb.TLSClientConfig,
-			Samples:   samples,
-			Options: lib.Options{
-				SystemTags: metrics.NewSystemTagSet(
-					metrics.TagName,
-					metrics.TagURL,
-				),
-				UserAgent: null.StringFrom("k6-test"),
-			},
-			BuiltinMetrics: metrics.RegisterBuiltinMetrics(
-				metrics.NewRegistry(),
-			),
-			Tags: lib.NewTagMap(nil),
-		}
+		testRuntime := modulestest.NewRuntime(t)
 
 		cwd, err := os.Getwd()
 		require.NoError(t, err)
@@ -105,15 +85,14 @@ func TestClient(t *testing.T) {
 				"file": fs,
 			},
 		}
+		testRuntime.InitContext(initEnv)
 
 		rt := goja.New()
 		rt.SetFieldNameMapper(common.FieldNameMapper{})
 
 		return testState{
-			rt:      rt,
+			Runtime: testRuntime,
 			httpBin: tb,
-			vuState: state,
-			env:     initEnv,
 			samples: samples,
 		}
 	}
@@ -716,31 +695,42 @@ func TestClient(t *testing.T) {
 
 			ts := setup(t)
 
-			ctx, cancel := context.WithCancel(context.Background())
-			defer cancel()
-			mvu := &modulestest.VU{
-				RuntimeField: ts.rt,
-				InitEnvField: ts.env,
-				CtxField:     ctx,
-			}
-
-			m, ok := New().NewModuleInstance(mvu).(*ModuleInstance)
+			m, ok := New().NewModuleInstance(ts.VU).(*ModuleInstance)
 			require.True(t, ok)
-			require.NoError(t, ts.rt.Set("grpc", m.Exports().Named))
+			require.NoError(t, ts.VU.Runtime().Set("grpc", m.Exports().Named))
 
 			// setup necessary environment if needed by a test
 			if tt.setup != nil {
 				tt.setup(ts.httpBin)
 			}
 
 			replace := func(code string) (goja.Value, error) {
-				return ts.rt.RunString(ts.httpBin.Replacer.Replace(code))
+				return ts.VU.Runtime().RunString(ts.httpBin.Replacer.Replace(code))
 			}
 
 			val, err := replace(tt.initString.code)
 			assertResponse(t, tt.initString, err, val, ts)
 
-			mvu.StateField = ts.vuState
+			root, err := lib.NewGroup("", nil)
+			require.NoError(t, err)
+			state := &lib.State{
+				Group:     root,
+				Dialer:    ts.httpBin.Dialer,
+				TLSConfig: ts.httpBin.TLSClientConfig,
+				Samples:   ts.samples,
+				Options: lib.Options{
+					SystemTags: metrics.NewSystemTagSet(
+						metrics.TagName,
+						metrics.TagURL,
+					),
+					UserAgent: null.StringFrom("k6-test"),
+				},
+				BuiltinMetrics: metrics.RegisterBuiltinMetrics(
+					metrics.NewRegistry(),
+				),
+				Tags: lib.NewTagMap(nil),
+			}
+			ts.MoveToVUContext(state)
 			val, err = replace(tt.vuString.code)
 			assertResponse(t, tt.vuString, err, val, ts)
 		})

js/modules/k6/ws/ws_test.go

@@ -96,24 +96,20 @@ func assertMetricEmittedCount(t *testing.T, metricName string, sampleContainers
 }
 
 type testState struct {
-	ctxPtr  *context.Context
-	rt      *goja.Runtime
+	*modulestest.Runtime
 	tb      *httpmultibin.HTTPMultiBin
-	state   *lib.State
 	samples chan metrics.SampleContainer
 }
 
 func newTestState(t testing.TB) testState {
 	tb := httpmultibin.NewHTTPMultiBin(t)
 
+	testRuntime := modulestest.NewRuntime(t)
+	samples := make(chan metrics.SampleContainer, 1000)
+
 	root, err := lib.NewGroup("", nil)
 	require.NoError(t, err)
 
-	rt := goja.New()
-	rt.SetFieldNameMapper(common.FieldNameMapper{})
-
-	samples := make(chan metrics.SampleContainer, 1000)
-
 	state := &lib.State{
 		Group:  root,
 		Dialer: tb.Dialer,
@@ -133,18 +129,13 @@ func newTestState(t testing.TB) testState {
 		Tags:           lib.NewTagMap(nil),
 	}
 
-	m := New().NewModuleInstance(&modulestest.VU{
-		CtxField:     tb.Context,
-		InitEnvField: &common.InitEnvironment{},
-		RuntimeField: rt,
-		StateField:   state,
-	})
-	require.NoError(t, rt.Set("ws", m.Exports().Default))
+	m := New().NewModuleInstance(testRuntime.VU)
+	require.NoError(t, testRuntime.VU.RuntimeField.Set("ws", m.Exports().Default))
+	testRuntime.MoveToVUContext(state)
 
 	return testState{
-		rt:      rt,
+		Runtime: testRuntime,
 		tb:      tb,
-		state:   state,
 		samples: samples,
 	}
 }
@@ -1049,7 +1040,7 @@ func TestCompression(t *testing.T) {
 			}
 		}))
 
-		_, err := ts.rt.RunString(sr(`
+		_, err := ts.VU.Runtime().RunString(sr(`
 		// if client supports compression, it has to send the header
 		// 'Sec-Websocket-Extensions:permessage-deflate; server_no_context_takeover; client_no_context_takeover' to server.
 		// if compression is negotiated successfully, server will reply with header
@@ -1131,7 +1122,7 @@ func TestCompression(t *testing.T) {
 					}
 				}))
 
-				_, err := ts.rt.RunString(sr(`
+				_, err := ts.VU.Runtime().RunString(sr(`
 					var res = ws.connect("WSBIN_URL/ws-compression-param", {"compression":"` + testCase.compression + `"}, function(socket){
 						socket.close()
 					});
@@ -1214,14 +1205,14 @@ func BenchmarkCompression(b *testing.B) {
 	b.ResetTimer()
 	b.Run("compression-enabled", func(b *testing.B) {
 		for i := 0; i < b.N; i++ {
-			if _, err := ts.rt.RunString(testCodes[0]); err != nil {
+			if _, err := ts.VU.Runtime().RunString(testCodes[0]); err != nil {
 				b.Error(err)
 			}
 		}
 	})
 	b.Run("compression-disabled", func(b *testing.B) {
 		for i := 0; i < b.N; i++ {
-			if _, err := ts.rt.RunString(testCodes[1]); err != nil {
+			if _, err := ts.VU.Runtime().RunString(testCodes[1]); err != nil {
 				b.Error(err)
 			}
 		}
@@ -1250,17 +1241,11 @@ func TestCookieJar(t *testing.T) {
 		}
 	}))
 
-	mii := &modulestest.VU{
-		RuntimeField: ts.rt,
-		InitEnvField: &common.InitEnvironment{Registry: metrics.NewRegistry()},
-		CtxField:     context.Background(),
-		StateField:   ts.state,
-	}
-	err := ts.rt.Set("http", httpModule.New().NewModuleInstance(mii).Exports().Default)
+	err := ts.VU.Runtime().Set("http", httpModule.New().NewModuleInstance(ts.VU).Exports().Default)
 	require.NoError(t, err)
-	ts.state.CookieJar, _ = cookiejar.New(nil)
+	ts.VU.State().CookieJar, _ = cookiejar.New(nil)
 
-	_, err = ts.rt.RunString(sr(`
+	_, err = ts.VU.Runtime().RunString(sr(`
 		var res = ws.connect("WSBIN_URL/ws-echo-someheader", function(socket){
 			socket.close()
 		})

js/modulestest/runtime.go

@@ -0,0 +1,51 @@
+// Package modulestest contains helpers to test js modules
+package modulestest
+
+import (
+	"context"
+	"testing"
+
+	"github.com/dop251/goja"
+	"go.k6.io/k6/js/common"
+	"go.k6.io/k6/js/eventloop"
+	"go.k6.io/k6/lib"
+)
+
+// Runtime is a helper struct that contains what is needed to run a (simple) module test
+type Runtime struct {
+	VU            *VU
+	EventLoop     *eventloop.EventLoop
+	CancelContext func()
+}
+
+// NewRuntime will create a new test runtime and will cancel the context on test/benchmark end
+func NewRuntime(t testing.TB) *Runtime {
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+	vu := &VU{
+		CtxField:     ctx,
+		RuntimeField: goja.New(),
+	}
+	vu.RuntimeField.SetFieldNameMapper(common.FieldNameMapper{})
+	eventloop := eventloop.New(vu)
+	vu.RegisterCallbackField = eventloop.RegisterCallback
+	result := &Runtime{
+		VU:            vu,
+		EventLoop:     eventloop,
+		CancelContext: cancel,
+	}
+	// let's cancel again in case it has changed
+	t.Cleanup(func() { result.CancelContext() })
+	return result
+}
+
+// InitContext will set the InitContext of the Runtime
+func (r *Runtime) InitContext(i *common.InitEnvironment) {
+	r.VU.InitEnvField = i
+}
+
+// MoveToVUContext will set the state and nil the InitEnv just as a real VU
+func (r *Runtime) MoveToVUContext(state *lib.State) {
+	r.VU.InitEnvField = nil
+	r.VU.StateField = state
+}