Use moveToContext only when needed in tests (#525)

Author: mstoykov

common/network_manager_test.go

@@ -57,6 +57,7 @@ func newTestNetworkManager(t *testing.T, k6opts k6lib.Options) (*NetworkManager,
 	}, nil)
 
 	vu := k6test.NewVU(t)
+	vu.MoveToVUContext()
 	st := vu.State()
 	st.Options = k6opts
 	logger := log.New(st.Logger, false, nil)

k6ext/k6test/vu.go

@@ -26,7 +26,8 @@ import (
 // RunLoop() override to call WaitOnRegistered()?
 type VU struct {
 	*k6modulestest.VU
-	Loop *k6eventloop.EventLoop
+	Loop      *k6eventloop.EventLoop
+	toBeState *k6lib.State
 }
 
 // ToGojaValue is a convenience method for converting any value to a goja value.
@@ -38,6 +39,13 @@ func (v *VU) RunLoop(fn func() error) error {
 	return v.Loop.Start(fn)
 }
 
+// MoveToVUContext moves the VU to VU context, adding a predefined k6 lib State and nilling the InitEnv
+// to simulate how that is done in the real k6.
+func (v *VU) MoveToVUContext() {
+	v.VU.StateField = v.toBeState
+	v.VU.InitEnvField = nil
+}
+
 // NewVU returns a mock k6 VU.
 func NewVU(tb testing.TB) *VU {
 	tb.Helper()
@@ -71,7 +79,6 @@ func NewVU(tb testing.TB) *VU {
 	testRT := k6modulestest.NewRuntime(tb)
 	ctx := k6ext.WithVU(testRT.VU.CtxField, testRT.VU)
 	testRT.VU.CtxField = ctx
-	testRT.MoveToVUContext(state)
 
-	return &VU{VU: testRT.VU, Loop: testRT.EventLoop}
+	return &VU{VU: testRT.VU, Loop: testRT.EventLoop, toBeState: state}
 }

tests/test_browser.go

@@ -33,7 +33,6 @@ import (
 	"github.com/grafana/xk6-browser/k6ext"
 	"github.com/grafana/xk6-browser/k6ext/k6test"
 
-	k6common "go.k6.io/k6/js/common"
 	k6http "go.k6.io/k6/js/modules/k6/http"
 	k6httpmultibin "go.k6.io/k6/lib/testutils/httpmultibin"
 	k6metrics "go.k6.io/k6/metrics"
@@ -105,43 +104,28 @@ func newTestBrowser(tb testing.TB, opts ...interface{}) *testBrowser {
 	k6m := k6ext.RegisterCustomMetrics(registry)
 	vu.CtxField = k6ext.WithCustomMetrics(vu.Context(), k6m)
 
+	v := chromium.NewBrowserType(vu)
+	bt, ok := v.(*chromium.BrowserType)
+	if !ok {
+		panic(fmt.Errorf("testBrowser: unexpected browser type %T", v))
+	}
+	vu.MoveToVUContext()
+	// enable the HTTP test server only when necessary
 	var (
-		state = vu.StateField
-		rt    = vu.RuntimeField
+		testServer *k6httpmultibin.HTTPMultiBin
+		state      = vu.StateField
+		rt         = vu.RuntimeField
+		lc         *logCache
 	)
 
-	// enable the HTTP test server only when necessary
-	var testServer *k6httpmultibin.HTTPMultiBin
+	if enableLogCache {
+		lc = attachLogCache(state.Logger)
+	}
 	if enableHTTPMultiBin {
 		testServer = k6httpmultibin.NewHTTPMultiBin(tb)
 		state.TLSConfig = testServer.TLSClientConfig
 		state.Transport = testServer.HTTPTransport
 	}
-
-	var lc *logCache
-	if enableLogCache {
-		lc = attachLogCache(state.Logger)
-	}
-
-	// launch the browser
-	initEnv := &k6common.InitEnvironment{
-		Logger:   state.Logger,
-		Registry: registry,
-	}
-	// NewBrowserType must be run from the init context, but since the recent
-	// validation introduced in k6/js/modulestest, InitEnvField and StateField
-	// must not be both set at the same time. So we nil StateField, create
-	// the BrowserType, and revert it before running Launch().
-	// See https://github.com/grafana/k6/pull/2602
-	vu.InitEnvField = initEnv
-	vu.StateField = nil
-	v := chromium.NewBrowserType(vu)
-	bt, ok := v.(*chromium.BrowserType)
-	if !ok {
-		panic(fmt.Errorf("testBrowser: unexpected browser type %T", v))
-	}
-	vu.InitEnvField = nil
-	vu.StateField = state
 	b := bt.Launch(rt.ToValue(launchOpts))
 	tb.Cleanup(func() {
 		select {