Refactor integration test browser

k6ext/k6test/vu.go

@@ -53,9 +53,9 @@ func (v *VU) AssertSamples(assertSample func(s k6metrics.Sample)) int {
 	return n
 }
 
-// WithSamplesListener is used to indicate we want to use a bidirectional channel
+// WithSamples is used to indicate we want to use a bidirectional channel
 // so that the test can read the metrics being emitted to the channel.
-type WithSamplesListener chan k6metrics.SampleContainer
+type WithSamples chan k6metrics.SampleContainer
 
 // NewVU returns a mock k6 VU.
 //
@@ -71,7 +71,7 @@ func NewVU(tb testing.TB, opts ...any) *VU {
 	)
 	for _, opt := range opts {
 		switch opt := opt.(type) {
-		case WithSamplesListener:
+		case WithSamples:
 			samples = opt
 		case env.LookupFunc:
 			lookupFunc = opt

tests/browser_test.go

@@ -45,7 +45,11 @@ func TestTmpDirCleanup(t *testing.T) {
 
 	const tmpDirPath = "./"
 
-	b := newTestBrowser(t, withSkipClose(), env.ConstLookup("TMPDIR", tmpDirPath))
+	b := newTestBrowser(
+		t,
+		withSkipClose(),
+		withEnvLookup(env.ConstLookup("TMPDIR", tmpDirPath)),
+	)
 	p := b.NewPage(nil)
 	err := p.Close(nil)
 	require.NoError(t, err)

tests/frame_test.go

@@ -81,7 +81,11 @@ func TestFrameNoPanicWithEmbeddedIFrame(t *testing.T) {
 	}
 
 	// run the browser in headfull mode.
-	tb := newTestBrowser(t, withFileServer(), env.ConstLookup(env.BrowserHeadless, "0"))
+	tb := newTestBrowser(
+		t,
+		withFileServer(),
+		withEnvLookup(env.ConstLookup(env.BrowserHeadless, "0")),
+	)
 
 	p := tb.NewPage(nil)
 	_, err := p.Goto(

tests/test_browser.go

@@ -47,111 +47,54 @@ type testBrowser struct {
 // It automatically closes it when `t` returns unless `withSkipClose` option is provided.
 //
 // The following opts are available to customize the testBrowser:
-//   - withHTTPServer: enables the HTTPMultiBin server.
+//   - withEnvLookup: provides a custom lookup function for environment variables.
 //   - withFileServer: enables the HTTPMultiBin server and serves the given files.
+//   - withHTTPServer: enables the HTTPMultiBin server.
 //   - withLogCache: enables the log cache.
-//   - withSamplesListener: provides a channel to receive the browser metrics.
-//   - env.LookupFunc: provides a custom lookup function for environment variables.
+//   - withSamples: provides a channel to receive the browser metrics.
 //   - withSkipClose: skips closing the browser when the test finishes.
-func newTestBrowser(tb testing.TB, opts ...any) *testBrowser {
+func newTestBrowser(tb testing.TB, opts ...func(*testBrowserOptions)) *testBrowser {
 	tb.Helper()
 
-	// set default options and then customize them
-	var (
-		enableHTTPMultiBin = false
-		enableFileServer   = false
-		enableLogCache     = false
-		skipClose          = false
-		samples            = make(chan k6metrics.SampleContainer, 1000)
-		// default lookup function is env.Lookup so that we can
-		// pass the environment variables while testing, i.e.: K6_BROWSER_LOG.
-		lookupFunc = env.Lookup
-	)
-	for _, opt := range opts {
-		switch opt := opt.(type) {
-		case httpServerOption:
-			enableHTTPMultiBin = true
-		case fileServerOption:
-			enableFileServer = true
-			enableHTTPMultiBin = true
-		case logCacheOption:
-			enableLogCache = true
-		case skipCloseOption:
-			skipClose = true
-		case withSamplesListener:
-			samples = opt
-		case env.LookupFunc:
-			lookupFunc = opt
-		}
-	}
-
-	vu := setupHTTPTestModuleInstance(tb, samples)
-
-	dummyCtx, cancel := context.WithCancel(vu.Context())
-	tb.Cleanup(cancel)
-	vu.CtxField = dummyCtx
-
-	registry := k6metrics.NewRegistry()
-	k6m := k6ext.RegisterCustomMetrics(registry)
-	vu.CtxField = k6ext.WithCustomMetrics(vu.Context(), k6m)
-	vu.InitEnvField.LookupEnv = lookupFunc
-
-	bt := chromium.NewBrowserType(vu)
-
-	// Delete the pre-init stage data.
-	vu.MoveToVUContext()
-
-	// enable the HTTP test server only when necessary
-	var (
-		testServer *k6httpmultibin.HTTPMultiBin
-		state      = vu.StateField
-		lc         *logCache
-	)
+	tbr := &testBrowser{t: tb}
+	tbopts := newTestBrowserOptions(opts...)
+	tbr.browserType, tbr.vu, tbr.cancel = newBrowserTypeWithVU(tb, tbopts)
+	tb.Cleanup(tbr.cancel)
 
-	if enableLogCache {
-		lc = attachLogCache(tb, state.Logger)
+	if tbopts.logCache {
+		tbr.logCache = attachLogCache(tb, tbr.vu.StateField.Logger)
 	}
-	if enableHTTPMultiBin {
-		testServer = k6httpmultibin.NewHTTPMultiBin(tb)
-		state.TLSConfig = testServer.TLSClientConfig
-		state.Transport = testServer.HTTPTransport
+	if tbopts.httpMultiBin {
+		tbr.http = k6httpmultibin.NewHTTPMultiBin(tb)
+		tbr.vu.StateField.TLSConfig = tbr.http.TLSClientConfig
+		tbr.vu.StateField.Transport = tbr.http.HTTPTransport
+	}
+	if tbopts.fileServer {
+		tbr = tbr.withFileServer()
 	}
 
-	b, pid, err := bt.Launch(dummyCtx)
+	b, pid, err := tbr.browserType.Launch(tbr.vu.Context())
 	if err != nil {
 		tb.Fatalf("testBrowser: %v", err)
 	}
 	cb, ok := b.(*common.Browser)
 	if !ok {
 		tb.Fatalf("testBrowser: unexpected browser %T", b)
 	}
-
+	tbr.Browser = cb
+	tbr.ctx = tbr.browserType.Ctx
+	tbr.pid = pid
+	tbr.wsURL = cb.WsURL()
 	tb.Cleanup(func() {
 		select {
-		case <-vu.Context().Done():
+		case <-tbr.vu.Context().Done():
 		default:
-			if !skipClose {
-				b.Close()
+			if !tbopts.skipClose {
+				cb.Close()
 			}
 		}
 	})
 
-	tbr := &testBrowser{
-		t:           tb,
-		ctx:         bt.Ctx,
-		http:        testServer,
-		vu:          vu,
-		logCache:    lc,
-		Browser:     cb,
-		browserType: bt,
-		pid:         pid,
-		wsURL:       cb.WsURL(),
-		cancel:      cancel,
-	}
-	if enableFileServer {
-		tbr = tbr.withFileServer()
-	}
-
 	return tbr
 }
 
@@ -169,38 +112,8 @@ func (b *testBrowser) NewPage(opts goja.Value) *common.Page {
 	return pp
 }
 
-// withHandler adds the given handler to the HTTP test server and makes it
-// accessible with the given pattern.
-func (b *testBrowser) withHandler(pattern string, handler http.HandlerFunc) *testBrowser {
-	b.t.Helper()
-
-	if b.http == nil {
-		b.t.Fatalf("You should enable HTTP test server, see: withHTTPServer option")
-	}
-	b.http.Mux.Handle(pattern, handler)
-	return b
-}
-
 const testBrowserStaticDir = "static"
 
-// withFileServer serves a file server using the HTTP test server that is
-// accessible via `testBrowserStaticDir` prefix.
-//
-// This method is for enabling the static file server after starting a test
-// browser. For early starting the file server see withFileServer function.
-func (b *testBrowser) withFileServer() *testBrowser {
-	b.t.Helper()
-
-	const (
-		slash = string(os.PathSeparator)
-		path  = slash + testBrowserStaticDir + slash
-	)
-
-	fs := http.FileServer(http.Dir(testBrowserStaticDir))
-
-	return b.withHandler(path, http.StripPrefix(path, fs).ServeHTTP)
-}
-
 // URL returns the listening HTTP test server's URL combined with the given path.
 func (b *testBrowser) URL(path string) string {
 	b.t.Helper()
@@ -341,23 +254,56 @@ func (b *testBrowser) awaitWithTimeout(timeout time.Duration, fn func() error) e
 	}
 }
 
-// httpServerOption is used to detect whether to enable the HTTP test
-// server.
-type httpServerOption struct{}
+// withFileServer serves a file server using the HTTP test server that is
+// accessible via `testBrowserStaticDir` prefix.
+//
+// This method is for enabling the static file server after starting a test
+// browser. For early starting the file server see withFileServer function.
+func (b *testBrowser) withFileServer() *testBrowser {
+	b.t.Helper()
+
+	const (
+		slash = string(os.PathSeparator)
+		path  = slash + testBrowserStaticDir + slash
+	)
+
+	fs := http.FileServer(http.Dir(testBrowserStaticDir))
+
+	return b.withHandler(path, http.StripPrefix(path, fs).ServeHTTP)
+}
+
+type testBrowserOptions struct {
+	fileServer   bool
+	logCache     bool
+	httpMultiBin bool
+	samples      chan k6metrics.SampleContainer
+	skipClose    bool
+	lookupFunc   env.LookupFunc
+}
 
-// withHTTPServer enables the HTTP test server.
+func newTestBrowserOptions(opts ...func(*testBrowserOptions)) *testBrowserOptions {
+	// default lookup function is env.Lookup so that we can
+	// pass the environment variables while testing, i.e.: K6_BROWSER_LOG.
+	tbo := &testBrowserOptions{
+		samples:    make(chan k6metrics.SampleContainer, 1000),
+		lookupFunc: env.Lookup,
+	}
+	for _, opt := range opts {
+		opt(tbo)
+	}
+
+	return tbo
+}
+
+// withEnvLookup sets the lookup function for environment variables.
 //
 // example:
 //
-//	b := TestBrowser(t, withHTTPServer())
-func withHTTPServer() httpServerOption {
-	return struct{}{}
+//	b := TestBrowser(t, withEnvLookup(env.ConstLookup(env.BrowserHeadless, "0")))
+func withEnvLookup(lookupFunc env.LookupFunc) func(*testBrowserOptions) {
+	return func(tb *testBrowserOptions) { tb.lookupFunc = lookupFunc }
 }
 
-// fileServerOption is used to detect whether enable the static file
-// server.
-type fileServerOption struct{}
-
 // withFileServer enables the HTTP test server and serves a file server
 // for static files.
 //
@@ -366,51 +312,84 @@ type fileServerOption struct{}
 // example:
 //
 //	b := TestBrowser(t, withFileServer())
-func withFileServer() fileServerOption {
-	return struct{}{}
+func withFileServer() func(tb *testBrowserOptions) {
+	return func(tb *testBrowserOptions) {
+		tb.httpMultiBin = true
+		tb.fileServer = true
+	}
+}
+
+// withHandler adds the given handler to the HTTP test server and makes it
+// accessible with the given pattern.
+func (b *testBrowser) withHandler(pattern string, handler http.HandlerFunc) *testBrowser {
+	b.t.Helper()
+
+	if b.http == nil {
+		b.t.Fatalf("You should enable HTTP test server, see: withHTTPServer option")
+	}
+	b.http.Mux.Handle(pattern, handler)
+	return b
 }
 
-// logCacheOption is used to detect whether to enable the log cache.
-type logCacheOption struct{}
+// withHTTPServer enables the HTTP test server.
+// It is used to detect whether to enable the HTTP test server.
+//
+// example:
+//
+//	b := TestBrowser(t, withHTTPServer())
+func withHTTPServer() func(tb *testBrowserOptions) {
+	return func(tb *testBrowserOptions) { tb.httpMultiBin = true }
+}
 
 // withLogCache enables the log cache.
 //
 // example:
 //
 //	b := TestBrowser(t, withLogCache())
-func withLogCache() logCacheOption {
-	return struct{}{}
+func withLogCache() func(tb *testBrowserOptions) {
+	return func(tb *testBrowserOptions) { tb.logCache = true }
 }
 
-// skipCloseOption is used to indicate that we shouldn't call Browser.Close() in
-// t.Cleanup(), since it will presumably be done by the test.
-type skipCloseOption struct{}
+// withSamples is used to indicate we want to use a bidirectional channel
+// so that the test can read the metrics being emitted to the channel.
+func withSamples(sc chan k6metrics.SampleContainer) func(tb *testBrowserOptions) {
+	return func(tb *testBrowserOptions) { tb.samples = sc }
+}
 
 // withSkipClose skips calling Browser.Close() in t.Cleanup().
+// It indicates that we shouldn't call Browser.Close() in
+// t.Cleanup(), since it will presumably be done by the test.
 //
 // example:
 //
 //	b := TestBrowser(t, withSkipClose())
-func withSkipClose() skipCloseOption {
-	return struct{}{}
+func withSkipClose() func(tb *testBrowserOptions) {
+	return func(tb *testBrowserOptions) { tb.skipClose = true }
 }
 
-// withSamplesListener is used to indicate we want to use a bidirectional channel
-// so that the test can read the metrics being emitted to the channel.
-type withSamplesListener chan k6metrics.SampleContainer
-
-func setupHTTPTestModuleInstance(tb testing.TB, samples chan k6metrics.SampleContainer) *k6test.VU {
+func newBrowserTypeWithVU(tb testing.TB, opts *testBrowserOptions) (
+	_ *chromium.BrowserType,
+	_ *k6test.VU,
+	cancel func(),
+) {
 	tb.Helper()
 
-	var (
-		vu   = k6test.NewVU(tb, k6test.WithSamplesListener(samples))
-		root = k6http.New()
+	// Prepare the VU.
+	vu := k6test.NewVU(tb, k6test.WithSamples(opts.samples))
+	mi, ok := k6http.New().NewModuleInstance(vu).(*k6http.ModuleInstance)
+	require.Truef(tb, ok, "want *k6http.ModuleInstance; got %T", mi)
+	require.NoError(tb, vu.Runtime().Set("http", mi.Exports().Default))
+	metricsCtx := k6ext.WithCustomMetrics(
+		vu.Context(),
+		k6ext.RegisterCustomMetrics(k6metrics.NewRegistry()),
 	)
+	ctx, cancel := context.WithCancel(metricsCtx)
+	vu.CtxField = ctx
+	vu.InitEnvField.LookupEnv = opts.lookupFunc
 
-	mi, ok := root.NewModuleInstance(vu).(*k6http.ModuleInstance)
-	require.True(tb, ok)
-
-	require.NoError(tb, vu.Runtime().Set("http", mi.Exports().Default))
+	bt := chromium.NewBrowserType(vu)
+	// Delete the pre-init stage data.
+	vu.MoveToVUContext()
 
-	return vu
+	return bt, vu, cancel
 }