Merge pull request #918 from grafana/refactor/822-k6-lookupenv

Refactor to k6 lookupEnv and abstract env. var. lookup

.github/workflows/e2e.yml

@@ -67,7 +67,7 @@ jobs:
       - name: Run E2E tests
         run: |
           set -x
-          export XK6_HEADLESS=true
+          export K6_BROWSER_HEADLESS=true
           for f in examples/*.js; do
             ./k6extension run "$f"
           done

.github/workflows/test.yml

@@ -46,7 +46,7 @@ jobs:
             args[1]="1"
             export GOMAXPROCS=1
           fi
-          export XK6_HEADLESS=true
+          export K6_BROWSER_HEADLESS=true
           go test "${args[@]}" -timeout 5m ./...
 
   test-tip:
@@ -75,7 +75,7 @@ jobs:
           go version
           export GOMAXPROCS=2
           args=("-p" "2" "-race")
-          export XK6_HEADLESS=true
+          export K6_BROWSER_HEADLESS=true
           go test "${args[@]}" -timeout 5m ./...
 
   test-current-cov:
@@ -104,7 +104,7 @@ jobs:
             args[1]="1"
             export GOMAXPROCS=1
           fi
-          export XK6_HEADLESS=true
+          export K6_BROWSER_HEADLESS=true
           echo "mode: set" > coverage.txt
           for pkg in $(go list ./... | grep -v vendor); do
               list=$(go list -test -f  '{{ join .Deps  "\n"}}' $pkg | grep github.com/grafana/xk6-browser | grep -v vendor || true)
@@ -156,5 +156,5 @@ jobs:
           go get go.k6.io/k6@master
           go mod tidy
           cat go.mod | grep go.k6.io/k6
-          export XK6_HEADLESS=true
+          export K6_BROWSER_HEADLESS=true
           go test "${args[@]}" -timeout 5m ./...

Dockerfile

@@ -11,6 +11,6 @@ RUN apt-get update && \
 
 COPY --from=builder /tmp/k6 /usr/bin/k6
 
-ENV XK6_HEADLESS=true
+ENV K6_BROWSER_HEADLESS=true
 
 ENTRYPOINT ["k6"]

browser/module.go

@@ -2,12 +2,15 @@
 package browser
 
 import (
-	"os"
+	"log"
+	"net/http"
+	_ "net/http/pprof" //nolint:gosec
 	"sync"
 
 	"github.com/dop251/goja"
 
 	"github.com/grafana/xk6-browser/common"
+	"github.com/grafana/xk6-browser/env"
 	"github.com/grafana/xk6-browser/k6ext"
 
 	k6modules "go.k6.io/k6/js/modules"
@@ -20,7 +23,7 @@ type (
 		PidRegistry     *pidRegistry
 		browserRegistry *browserRegistry
 		remoteRegistry  *remoteRegistry
-		once            *sync.Once
+		initOnce        *sync.Once
 	}
 
 	// JSModule exposes the properties available to the JS script.
@@ -46,27 +49,21 @@ func New() *RootModule {
 	return &RootModule{
 		PidRegistry:     &pidRegistry{},
 		browserRegistry: &browserRegistry{},
-		once:            &sync.Once{},
+		initOnce:        &sync.Once{},
 	}
 }
 
 // NewModuleInstance implements the k6modules.Module interface to return
 // a new instance for each VU.
 func (m *RootModule) NewModuleInstance(vu k6modules.VU) k6modules.Instance {
-	m.once.Do(func() {
-		// remoteRegistry should only be initialized once as it is
-		// used globally across the whole test run and not just the
-		// current vu. Since newRemoteRegistry can fail with an error,
-		// we've had to place it here so that if an error occurs a
-		// panic can be initiated and safely handled by k6.
-		rr, err := newRemoteRegistry(os.LookupEnv)
-		if err != nil {
-			ctx := k6ext.WithVU(vu.Context(), vu)
-			k6ext.Panic(ctx, "failed to create remote registry: %v", err.Error())
-		}
-		m.remoteRegistry = rr
+	// initialization should be done once per module as it initializes
+	// globally used values across the whole test run and not just the
+	// current VU. Since initialization can fail with an error,
+	// we've had to place it here so that if an error occurs a
+	// panic can be initiated and safely handled by k6.
+	m.initOnce.Do(func() {
+		m.initialize(vu)
 	})
-
 	return &ModuleInstance{
 		mod: &JSModule{
 			Browser: mapBrowserToGoja(moduleVU{
@@ -85,3 +82,25 @@ func (m *RootModule) NewModuleInstance(vu k6modules.VU) k6modules.Instance {
 func (mi *ModuleInstance) Exports() k6modules.Exports {
 	return k6modules.Exports{Default: mi.mod}
 }
+
+// initialize initializes the module instance with a new remote registry
+// and debug server, etc.
+func (m *RootModule) initialize(vu k6modules.VU) {
+	var (
+		err     error
+		initEnv = vu.InitEnv()
+	)
+	m.remoteRegistry, err = newRemoteRegistry(initEnv.LookupEnv)
+	if err != nil {
+		k6ext.Abort(vu.Context(), "failed to create remote registry: %v", err)
+	}
+	if _, ok := initEnv.LookupEnv(env.EnableProfiling); ok {
+		go startDebugServer()
+	}
+}
+
+func startDebugServer() {
+	log.Println("Starting http debug server", env.ProfilingServerAddr)
+	log.Println(http.ListenAndServe(env.ProfilingServerAddr, nil)) //nolint:gosec
+	// no linted because we don't need to set timeouts for the debug server.
+}

browser/module_test.go

@@ -7,6 +7,8 @@ import (
 	"github.com/dop251/goja"
 	"github.com/stretchr/testify/require"
 
+	"github.com/grafana/xk6-browser/env"
+
 	k6common "go.k6.io/k6/js/common"
 	k6modulestest "go.k6.io/k6/js/modulestest"
 	k6lib "go.k6.io/k6/lib"
@@ -23,7 +25,8 @@ func TestModuleNew(t *testing.T) {
 		RuntimeField: goja.New(),
 		InitEnvField: &k6common.InitEnvironment{
 			TestPreInitState: &k6lib.TestPreInitState{
-				Registry: k6metrics.NewRegistry(),
+				Registry:  k6metrics.NewRegistry(),
+				LookupEnv: env.EmptyLookup,
 			},
 		},
 		CtxField: context.Background(),

browser/registry.go

@@ -70,7 +70,7 @@ func newRemoteRegistry(envLookup env.LookupFunc) (*remoteRegistry, error) {
 }
 
 func checkForBrowserWSURLs(envLookup env.LookupFunc) (bool, []string) {
-	wsURL, isRemote := envLookup("K6_BROWSER_WS_URL")
+	wsURL, isRemote := envLookup(env.WebSocketURLs)
 	if !isRemote {
 		return false, nil
 	}
@@ -93,7 +93,7 @@ func checkForBrowserWSURLs(envLookup env.LookupFunc) (bool, []string) {
 // checkForScenarios will parse the K6_INSTANCE_SCENARIOS env var if
 // it has been defined.
 func checkForScenarios(envLookup env.LookupFunc) (bool, []string, error) {
-	scenariosJSON, isRemote := envLookup("K6_INSTANCE_SCENARIOS")
+	scenariosJSON, isRemote := envLookup(env.InstanceScenarios)
 	if !isRemote {
 		return false, nil, nil
 	}

browser/registry_test.go

@@ -2,13 +2,14 @@ package browser
 
 import (
 	"errors"
-	"os"
 	"strconv"
 	"sync"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
+	"github.com/grafana/xk6-browser/env"
 )
 
 func TestPidRegistry(t *testing.T) {
@@ -34,118 +35,125 @@ func TestPidRegistry(t *testing.T) {
 }
 
 func TestIsRemoteBrowser(t *testing.T) {
+	t.Parallel()
+
 	testCases := []struct {
-		name           string
-		expIsRemote    bool
-		expValidWSURLs []string
-		envVarName     string
-		envVarValue    string
-		expErr         error
+		name                    string
+		envVarName, envVarValue string
+		expIsRemote             bool
+		expValidWSURLs          []string
+		expErr                  error
 	}{
 		{
 			name:        "browser is not remote",
-			expIsRemote: false,
 			envVarName:  "FOO",
 			envVarValue: "BAR",
+			expIsRemote: false,
 		},
 		{
 			name:           "single WS URL",
+			envVarName:     env.WebSocketURLs,
+			envVarValue:    "WS_URL",
 			expIsRemote:    true,
 			expValidWSURLs: []string{"WS_URL"},
-			envVarName:     "K6_BROWSER_WS_URL",
-			envVarValue:    "WS_URL",
 		},
 		{
 			name:           "multiple WS URL",
+			envVarName:     env.WebSocketURLs,
+			envVarValue:    "WS_URL_1,WS_URL_2,WS_URL_3",
 			expIsRemote:    true,
 			expValidWSURLs: []string{"WS_URL_1", "WS_URL_2", "WS_URL_3"},
-			envVarName:     "K6_BROWSER_WS_URL",
-			envVarValue:    "WS_URL_1,WS_URL_2,WS_URL_3",
 		},
 		{
 			name:           "ending comma is handled",
+			envVarName:     env.WebSocketURLs,
+			envVarValue:    "WS_URL_1,WS_URL_2,",
 			expIsRemote:    true,
 			expValidWSURLs: []string{"WS_URL_1", "WS_URL_2"},
-			envVarName:     "K6_BROWSER_WS_URL",
-			envVarValue:    "WS_URL_1,WS_URL_2,",
 		},
 		{
 			name:           "void string does not panic",
+			envVarName:     env.WebSocketURLs,
+			envVarValue:    "",
 			expIsRemote:    true,
 			expValidWSURLs: []string{""},
-			envVarName:     "K6_BROWSER_WS_URL",
-			envVarValue:    "",
 		},
 		{
 			name:           "comma does not panic",
+			envVarName:     env.WebSocketURLs,
+			envVarValue:    ",",
 			expIsRemote:    true,
 			expValidWSURLs: []string{""},
-			envVarName:     "K6_BROWSER_WS_URL",
-			envVarValue:    ",",
 		},
 		{
 			name:           "read a single scenario with a single ws url",
+			envVarName:     env.InstanceScenarios,
+			envVarValue:    `[{"id": "one","browsers": [{ "handle": "WS_URL_1" }]}]`,
 			expIsRemote:    true,
 			expValidWSURLs: []string{"WS_URL_1"},
-			envVarName:     "K6_INSTANCE_SCENARIOS",
-			envVarValue:    `[{"id": "one","browsers": [{ "handle": "WS_URL_1" }]}]`,
 		},
 		{
 			name:           "read a single scenario with a two ws urls",
+			envVarName:     env.InstanceScenarios,
+			envVarValue:    `[{"id": "one","browsers": [{"handle": "WS_URL_1"}, {"handle": "WS_URL_2"}]}]`,
 			expIsRemote:    true,
 			expValidWSURLs: []string{"WS_URL_1", "WS_URL_2"},
-			envVarName:     "K6_INSTANCE_SCENARIOS",
-			envVarValue:    `[{"id": "one","browsers": [{"handle": "WS_URL_1"}, {"handle": "WS_URL_2"}]}]`,
 		},
 		{
-			name:           "read two scenarios with multiple ws urls",
+			name:       "read two scenarios with multiple ws urls",
+			envVarName: env.InstanceScenarios,
+			envVarValue: `[
+				{"id": "one","browsers": [{"handle": "WS_URL_1"}, {"handle": "WS_URL_2"}]},
+				{"id": "two","browsers": [{"handle": "WS_URL_3"}, {"handle": "WS_URL_4"}]}
+			]`,
 			expIsRemote:    true,
 			expValidWSURLs: []string{"WS_URL_1", "WS_URL_2", "WS_URL_3", "WS_URL_4"},
-			envVarName:     "K6_INSTANCE_SCENARIOS",
-			envVarValue: `[{"id": "one","browsers": [{"handle": "WS_URL_1"}, {"handle": "WS_URL_2"}]},
-			{"id": "two","browsers": [{"handle": "WS_URL_3"}, {"handle": "WS_URL_4"}]}]`,
 		},
 		{
 			name:           "read scenarios without any ws urls",
+			envVarName:     env.InstanceScenarios,
+			envVarValue:    `[{"id": "one","browsers": [{}]}]`,
 			expIsRemote:    false,
 			expValidWSURLs: []string{""},
-			envVarName:     "K6_INSTANCE_SCENARIOS",
-			envVarValue:    `[{"id": "one","browsers": [{}]}]`,
 		},
 		{
 			name:           "read scenarios without any browser objects",
+			envVarName:     env.InstanceScenarios,
+			envVarValue:    `[{"id": "one"}]`,
 			expIsRemote:    false,
 			expValidWSURLs: []string{""},
-			envVarName:     "K6_INSTANCE_SCENARIOS",
-			envVarValue:    `[{"id": "one"}]`,
 		},
 		{
 			name:        "read empty scenarios",
-			expErr:      errors.New("parsing K6_INSTANCE_SCENARIOS: unexpected end of JSON input"),
-			envVarName:  "K6_INSTANCE_SCENARIOS",
+			envVarName:  env.InstanceScenarios,
 			envVarValue: ``,
+			expErr:      errors.New("parsing K6_INSTANCE_SCENARIOS: unexpected end of JSON input"),
 		},
 	}
-
 	for _, tc := range testCases {
+		tc := tc
 		t.Run(tc.name, func(t *testing.T) {
-			// the real environment  variable we receive needs quoting.
-			// this makes it as close to the real implementation as possible.
-			v := tc.envVarValue
-			if tc.envVarName == "K6_INSTANCE_SCENARIOS" {
-				v = strconv.Quote(v)
+			t.Parallel()
+
+			lookup := func(key string) (string, bool) {
+				v := tc.envVarValue
+				if tc.envVarName == "K6_INSTANCE_SCENARIOS" {
+					v = strconv.Quote(v)
+				}
+				if key == tc.envVarName {
+					return v, true
+				}
+				return "", false
 			}
-			t.Setenv(tc.envVarName, v)
 
-			rr, err := newRemoteRegistry(os.LookupEnv)
+			rr, err := newRemoteRegistry(lookup)
 			if tc.expErr != nil {
 				assert.Error(t, tc.expErr, err)
 				return
 			}
 			assert.NoError(t, err)
 
 			wsURL, isRemote := rr.isRemoteBrowser()
-
 			require.Equal(t, tc.expIsRemote, isRemote)
 			if isRemote {
 				require.Contains(t, tc.expValidWSURLs, wsURL)
@@ -154,10 +162,18 @@ func TestIsRemoteBrowser(t *testing.T) {
 	}
 
 	t.Run("K6_INSTANCE_SCENARIOS should override K6_BROWSER_WS_URL", func(t *testing.T) {
-		t.Setenv("K6_BROWSER_WS_URL", "WS_URL_1")
-		t.Setenv("K6_INSTANCE_SCENARIOS", strconv.Quote(`[{"id": "one","browsers": [{ "handle": "WS_URL_2" }]}]`))
+		lookup := func(key string) (string, bool) {
+			switch key {
+			case env.WebSocketURLs:
+				return "WS_URL_1", true
+			case env.InstanceScenarios:
+				return strconv.Quote(`[{"id": "one","browsers": [{ "handle": "WS_URL_2" }]}]`), true
+			default:
+				return "", false
+			}
+		}
 
-		rr, err := newRemoteRegistry(os.LookupEnv)
+		rr, err := newRemoteRegistry(lookup)
 		assert.NoError(t, err)
 
 		wsURL, isRemote := rr.isRemoteBrowser()