Refactor integration test browser

tests/test_browser.go

@@ -29,19 +29,31 @@ const testBrowserStaticDir = "static"
 type testBrowser struct {
 	t testing.TB
 
-	ctx      context.Context
-	http     *k6httpmultibin.HTTPMultiBin
-	vu       *k6test.VU
-	logCache *logCache
-
-	pid   int // the browser process ID
-	wsURL string
+	ctx    context.Context
+	cancel context.CancelFunc
+	vu     *k6test.VU
 
 	browserType *chromium.BrowserType
+	pid         int // the browser process ID
+	wsURL       string
 
 	*common.Browser
 
-	cancel context.CancelFunc
+	// isBrowserTypeInitialized is true if the browser type has been
+	// initialized with a VU. Some options can only be used in the
+	// post-init stage and require the browser type to be initialized.
+	isBrowserTypeInitialized bool
+
+	// http is set by the withHTTPServer option.
+	http *k6httpmultibin.HTTPMultiBin
+	// logCache is set by the withLogCache option.
+	logCache *logCache
+	// lookupFunc is set by the withEnvLookup option.
+	lookupFunc env.LookupFunc
+	// samples is set by the withSamples option.
+	samples chan k6metrics.SampleContainer
+	// skipClose is set by the withSkipClose option.
+	skipClose bool
 }
 
 // newTestBrowser configures and launches a new chrome browser.
@@ -55,25 +67,16 @@ type testBrowser struct {
 //   - withLogCache: enables the log cache.
 //   - withSamples: provides a channel to receive the browser metrics.
 //   - withSkipClose: skips closing the browser when the test finishes.
-func newTestBrowser(tb testing.TB, opts ...func(*testBrowserOptions)) *testBrowser {
+func newTestBrowser(tb testing.TB, opts ...func(*testBrowser)) *testBrowser {
 	tb.Helper()
 
 	tbr := &testBrowser{t: tb}
-	tbopts := newTestBrowserOptions(opts...)
-	tbr.browserType, tbr.vu, tbr.cancel = newBrowserTypeWithVU(tb, tbopts)
+	tbr.applyDefaultOptions()
+	tbr.applyOptions(opts...) // apply pre-init stage options.
+	tbr.browserType, tbr.vu, tbr.cancel = newBrowserTypeWithVU(tb, tbr)
 	tb.Cleanup(tbr.cancel)
-
-	if tbopts.logCache {
-		tbr.logCache = attachLogCache(tb, tbr.vu.StateField.Logger)
-	}
-	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()
-	}
+	tbr.isBrowserTypeInitialized = true // some option require the browser type to be initialized.
+	tbr.applyOptions(opts...)           // apply post-init stage options.
 
 	b, pid, err := tbr.browserType.Launch(tbr.vu.Context())
 	if err != nil {
@@ -91,7 +94,7 @@ func newTestBrowser(tb testing.TB, opts ...func(*testBrowserOptions)) *testBrows
 		select {
 		case <-tbr.vu.Context().Done():
 		default:
-			if !tbopts.skipClose {
+			if !tbr.skipClose {
 				cb.Close()
 			}
 		}
@@ -100,6 +103,160 @@ func newTestBrowser(tb testing.TB, opts ...func(*testBrowserOptions)) *testBrows
 	return tbr
 }
 
+// newBrowserTypeWithVU creates a new browser type with a VU.
+func newBrowserTypeWithVU(tb testing.TB, tbr *testBrowser) (
+	_ *chromium.BrowserType,
+	_ *k6test.VU,
+	cancel func(),
+) {
+	tb.Helper()
+
+	vu := k6test.NewVU(tb, k6test.WithSamples(tbr.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 = tbr.lookupFunc
+
+	bt := chromium.NewBrowserType(vu)
+	vu.RestoreVUState()
+
+	return bt, vu, cancel
+}
+
+// applyDefaultOptions applies the default options for the testBrowser.
+func (b *testBrowser) applyDefaultOptions() {
+	b.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.
+	b.lookupFunc = env.Lookup
+}
+
+// applyOptions applies the given options to the testBrowser.
+func (b *testBrowser) applyOptions(opts ...func(*testBrowser)) {
+	for _, opt := range opts {
+		opt(b)
+	}
+}
+
+// withEnvLookup sets the lookup function for environment variables.
+//
+// example:
+//
+//	b := TestBrowser(t, withEnvLookup(env.ConstLookup(env.BrowserHeadless, "0")))
+func withEnvLookup(lookupFunc env.LookupFunc) func(*testBrowser) {
+	return func(tb *testBrowser) { tb.lookupFunc = lookupFunc }
+}
+
+// withFileServer enables the HTTP test server and serves a file server
+// for static files.
+//
+// see: WithFileServer
+//
+// example:
+//
+//	b := TestBrowser(t, withFileServer())
+func withFileServer() func(*testBrowser) {
+	return func(tb *testBrowser) {
+		if !tb.isBrowserTypeInitialized {
+			return
+		}
+		if tb.http == nil {
+			// file server needs HTTP server.
+			apply := withHTTPServer()
+			apply(tb)
+		}
+		_ = tb.withFileServer()
+	}
+}
+
+// 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)
+}
+
+// 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
+}
+
+// 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(*testBrowser) {
+	return func(tb *testBrowser) {
+		if !tb.isBrowserTypeInitialized {
+			return
+		}
+		if tb.http != nil {
+			// already initialized.
+			return
+		}
+		tb.http = k6httpmultibin.NewHTTPMultiBin(tb.t)
+		tb.vu.StateField.TLSConfig = tb.http.TLSClientConfig
+		tb.vu.StateField.Transport = tb.http.HTTPTransport
+	}
+}
+
+// withLogCache enables the log cache.
+//
+// example:
+//
+//	b := TestBrowser(t, withLogCache())
+func withLogCache() func(*testBrowser) {
+	return func(tb *testBrowser) {
+		if !tb.isBrowserTypeInitialized {
+			return
+		}
+		tb.logCache = attachLogCache(tb.t, tb.vu.StateField.Logger)
+	}
+}
+
+// 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(*testBrowser) {
+	return func(tb *testBrowser) { 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() func(*testBrowser) {
+	return func(tb *testBrowser) { tb.skipClose = true }
+}
+
 // NewPage is a wrapper around Browser.NewPage that fails the test if an
 // error occurs. Added this helper to avoid boilerplate code in tests.
 func (b *testBrowser) NewPage(opts goja.Value) *common.Page {
@@ -218,151 +375,3 @@ func (b *testBrowser) awaitWithTimeout(timeout time.Duration, fn func() error) e
 		return fmt.Errorf("test timed out after %s", timeout)
 	}
 }
-
-// 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)
-}
-
-// testBrowserOptions is a helper for creating testBrowser options.
-type testBrowserOptions struct {
-	fileServer   bool
-	logCache     bool
-	httpMultiBin bool
-	samples      chan k6metrics.SampleContainer
-	skipClose    bool
-	lookupFunc   env.LookupFunc
-}
-
-// newTestBrowserOptions creates a new testBrowserOptions with the given options.
-// call apply to reapply the options.
-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,
-	}
-	tbo.apply(opts...)
-
-	return tbo
-}
-
-// apply applies the given options to the testBrowserOptions.
-func (tbo *testBrowserOptions) apply(opts ...func(*testBrowserOptions)) {
-	for _, opt := range opts {
-		opt(tbo)
-	}
-}
-
-// withEnvLookup sets the lookup function for environment variables.
-//
-// example:
-//
-//	b := TestBrowser(t, withEnvLookup(env.ConstLookup(env.BrowserHeadless, "0")))
-func withEnvLookup(lookupFunc env.LookupFunc) func(*testBrowserOptions) {
-	return func(tb *testBrowserOptions) { tb.lookupFunc = lookupFunc }
-}
-
-// withFileServer enables the HTTP test server and serves a file server
-// for static files.
-//
-// see: WithFileServer
-//
-// example:
-//
-//	b := TestBrowser(t, withFileServer())
-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
-}
-
-// 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() func(tb *testBrowserOptions) {
-	return func(tb *testBrowserOptions) { tb.logCache = true }
-}
-
-// 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() func(tb *testBrowserOptions) {
-	return func(tb *testBrowserOptions) { tb.skipClose = true }
-}
-
-func newBrowserTypeWithVU(tb testing.TB, opts *testBrowserOptions) (
-	_ *chromium.BrowserType,
-	_ *k6test.VU,
-	cancel func(),
-) {
-	tb.Helper()
-
-	// 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
-
-	bt := chromium.NewBrowserType(vu)
-	// Delete the pre-init stage data.
-	vu.RestoreVUState()
-
-	return bt, vu, cancel
-}