Add lifecycle event validation to WaitForLoadState

Author: Ivan Mirić

common/frame.go

@@ -1583,12 +1583,9 @@ func (f *Frame) WaitForLoadState(state string, opts goja.Value) {
 		k6Throw(f.ctx, "cannot parse waitForLoadState options: %v", err)
 	}
 
-	waitUntil := LifecycleEventLoad
-	switch state {
-	case "domcontentloaded":
-		waitUntil = LifecycleEventDOMContentLoad
-	case "networkidle":
-		waitUntil = LifecycleEventNetworkIdle
+	var waitUntil LifecycleEvent
+	if err = waitUntil.UnmarshalText([]byte(state)); err != nil {
+		k6Throw(f.ctx, "waitForLoadState error: %v", err)
 	}
 
 	if f.hasLifecycleEventFired(waitUntil) {

common/types.go

@@ -26,6 +26,8 @@ import (
 	"encoding/json"
 	"fmt"
 	"math"
+	"sort"
+	"strings"
 
 	"github.com/dop251/goja"
 
@@ -278,6 +280,47 @@ func (l *LifecycleEvent) UnmarshalJSON(b []byte) error {
 	return nil
 }
 
+// MarshalText returns the string representation of the enum value.
+// It returns an error if the enum value is invalid.
+func (l *LifecycleEvent) MarshalText() ([]byte, error) {
+	if l == nil {
+		return []byte(""), nil
+	}
+	var (
+		ok bool
+		s  string
+	)
+	if s, ok = lifecycleEventToString[*l]; !ok {
+		return nil, fmt.Errorf("invalid lifecycle event: %v", int(*l))
+	}
+
+	return []byte(s), nil
+}
+
+// UnmarshalText unmarshals a text representation to the enum value.
+// It returns an error if given a wrong value.
+func (l *LifecycleEvent) UnmarshalText(text []byte) error {
+	var (
+		ok  bool
+		val = string(text)
+	)
+
+	if *l, ok = lifecycleEventToID[val]; !ok {
+		valid := make([]string, 0, len(lifecycleEventToID))
+		for k := range lifecycleEventToID {
+			valid = append(valid, k)
+		}
+		sort.Slice(valid, func(i, j int) bool {
+			return lifecycleEventToID[valid[j]] > lifecycleEventToID[valid[i]]
+		})
+		return fmt.Errorf(
+			"invalid lifecycle event: %q; must be one of: %s",
+			val, strings.Join(valid, ", "))
+	}
+
+	return nil
+}
+
 type MediaType string
 
 const (

common/types_test.go

@@ -4,6 +4,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 // See: Issue #183 for details.
@@ -60,3 +61,71 @@ func TestViewportCalculateInset(t *testing.T) {
 		})
 	}
 }
+
+func TestLifecycleEventMarshalText(t *testing.T) {
+	t.Parallel()
+
+	t.Run("ok/nil", func(t *testing.T) {
+		t.Parallel()
+
+		var evt *LifecycleEvent
+		m, err := evt.MarshalText()
+		require.NoError(t, err)
+		assert.Equal(t, []byte(""), m)
+	})
+
+	t.Run("err/invalid", func(t *testing.T) {
+		t.Parallel()
+
+		evt := LifecycleEvent(-1)
+		_, err := evt.MarshalText()
+		require.EqualError(t, err, "invalid lifecycle event: -1")
+	})
+}
+
+func TestLifecycleEventMarshalTextRound(t *testing.T) {
+	t.Parallel()
+
+	evt := LifecycleEventLoad
+	m, err := evt.MarshalText()
+	require.NoError(t, err)
+	assert.Equal(t, []byte("load"), m)
+
+	var evt2 LifecycleEvent
+	err = evt2.UnmarshalText(m)
+	require.NoError(t, err)
+	assert.Equal(t, evt, evt2)
+}
+
+func TestLifecycleEventUnmarshalText(t *testing.T) {
+	t.Parallel()
+
+	t.Run("ok", func(t *testing.T) {
+		t.Parallel()
+
+		var evt LifecycleEvent
+		err := evt.UnmarshalText([]byte("load"))
+		require.NoError(t, err)
+		assert.Equal(t, LifecycleEventLoad, evt)
+	})
+
+	t.Run("err/invalid", func(t *testing.T) {
+		t.Parallel()
+
+		var evt LifecycleEvent
+		err := evt.UnmarshalText([]byte("none"))
+		require.EqualError(t, err,
+			`invalid lifecycle event: "none"; `+
+				`must be one of: load, domcontentloaded, networkidle`)
+	})
+
+	t.Run("err/invalid_empty", func(t *testing.T) {
+		t.Parallel()
+
+		var evt LifecycleEvent
+		err := evt.UnmarshalText([]byte(""))
+		require.EqualError(t, err,
+			`invalid lifecycle event: ""; `+
+				`must be one of: load, domcontentloaded, networkidle`)
+	})
+}

tests/page_test.go

@@ -629,6 +629,26 @@ func TestPageWaitForFunction(t *testing.T) {
 	})
 }
 
+func TestPageWaitForLoadState(t *testing.T) {
+	t.Parallel()
+
+	// TODO: Add happy path tests once WaitForLoadState is not racy.
+
+	t.Run("err_wrong_event", func(t *testing.T) {
+		t.Parallel()
+
+		defer func() {
+			assertPanicErrorContains(t, recover(),
+				`invalid lifecycle event: "none"; `+
+					`must be one of: load, domcontentloaded, networkidle`)
+		}()
+
+		p := newTestBrowser(t).NewPage(nil)
+		p.WaitForLoadState("none", nil)
+		t.Error("did not panic")
+	})
+}
+
 // See: The issue #187 for details.
 func TestPageWaitForNavigationShouldNotPanic(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
@@ -645,8 +665,7 @@ func assertPanicErrorContains(t *testing.T, err interface{}, expErrMsg string) {
 	require.IsType(t, &goja.Object{}, err)
 	gotObj, _ := err.(*goja.Object)
 	got := gotObj.Export()
-	expErr := fmt.Errorf("%w", errors.New(expErrMsg))
-	require.IsType(t, expErr, got)
+	expErr := errors.New(expErrMsg)
 	gotErr, ok := got.(error)
 	require.True(t, ok)
 	assert.Contains(t, gotErr.Error(), expErr.Error())