Fix an issue with WaitForLoadState

This fixes an issue when we actually have to wait for a lifecycle
event to be received from the browser since we've not received one
yet and so we can't rely on the frame's state.

The logic to wait for the lifecycle event has been copied from
FrameManager.NavigateFrame.

Author: ankur22

common/frame.go

@@ -1818,6 +1818,7 @@ func (f *Frame) WaitForFunction(fn goja.Value, opts goja.Value, jsArgs ...goja.V
 }
 
 // WaitForLoadState waits for the given load state to be reached.
+// This will unblock if that lifecycle event has already been received.
 func (f *Frame) WaitForLoadState(state string, opts goja.Value) {
 	f.log.Debugf("Frame:WaitForLoadState", "fid:%s furl:%q state:%s", f.ID(), f.URL(), state)
 	defer f.log.Debugf("Frame:WaitForLoadState:return", "fid:%s furl:%q state:%s", f.ID(), f.URL(), state)
@@ -1828,21 +1829,33 @@ func (f *Frame) WaitForLoadState(state string, opts goja.Value) {
 		k6ext.Panic(f.ctx, "parsing waitForLoadState %q options: %v", state, err)
 	}
 
+	timeoutCtx, timeoutCancel := context.WithTimeout(f.ctx, parsedOpts.Timeout)
+	defer timeoutCancel()
+
 	waitUntil := LifecycleEventLoad
 	if state != "" {
 		if err = waitUntil.UnmarshalText([]byte(state)); err != nil {
 			k6ext.Panic(f.ctx, "waiting for load state: %v", err)
 		}
 	}
 
+	lifecycleEvtCh, lifecycleEvtCancel := createWaitForEventPredicateHandler(
+		timeoutCtx, f, []string{EventFrameAddLifecycle},
+		func(data any) bool {
+			if le, ok := data.(LifecycleEvent); ok {
+				return le == waitUntil
+			}
+			return false
+		})
+	defer lifecycleEvtCancel()
+
 	if f.hasLifecycleEventFired(waitUntil) {
 		return
 	}
 
-	_, err = waitForEvent(f.ctx, f, []string{EventFrameAddLifecycle}, func(data any) bool {
-		return data.(LifecycleEvent) == waitUntil
-	}, parsedOpts.Timeout)
-	if err != nil {
+	select {
+	case <-lifecycleEvtCh:
+	case <-timeoutCtx.Done():
 		k6ext.Panic(f.ctx, "waiting for load state %q: %v", state, err)
 	}
 }

tests/lifecycle_wait_test.go

@@ -174,6 +174,55 @@ func TestLifecycleWaitForLoadStateNetworkIdle(t *testing.T) {
 	})
 }
 
+func TestLifecycleWaitForLoadStateDOMContentLoadedThenNetworkIdle(t *testing.T) {
+	// Test description
+	//
+	// 1. goto /home and wait for the domcontentloaded lifecycle event.
+	// 2. use WaitForLoadState with networkidle to now wait for the
+	//    lifecycle event from the browser.
+	//
+	// Success criteria: We want to quickly move to calling WaitForLoadState
+	//                   so that we block until a networkidle lifecycle
+	//                   event is received from the browser.
+
+	t.Parallel()
+
+	tb := newTestBrowser(t, withFileServer())
+	p := tb.NewPage(nil)
+	tb.withHandler("/home", func(w http.ResponseWriter, r *http.Request) {
+		http.Redirect(w, r, tb.staticURL("wait_for_nav_lifecycle.html"), http.StatusMovedPermanently)
+	})
+
+	var counter int64
+	var counterMu sync.Mutex
+	tb.withHandler("/ping", func(w http.ResponseWriter, _ *http.Request) {
+		counterMu.Lock()
+		defer counterMu.Unlock()
+
+		time.Sleep(time.Millisecond * 100)
+
+		counter++
+		fmt.Fprintf(w, "pong %d", counter)
+	})
+
+	tb.withHandler("/ping.js", func(w http.ResponseWriter, _ *http.Request) {
+		fmt.Fprintf(w, `
+				var pingJSTextOutput = document.getElementById("pingJSText");
+				pingJSTextOutput.innerText = "ping.js loaded from server";
+			`)
+	})
+
+	assertHome(t, tb, p, common.LifecycleEventDOMContentLoad, func() {
+		p.WaitForLoadState(common.LifecycleEventNetworkIdle.String(), nil)
+
+		result := p.TextContent("#pingRequestText", nil)
+		assert.EqualValues(t, "Waiting... pong 10 - for loop complete", result)
+
+		result = p.TextContent("#pingJSText", nil)
+		assert.EqualValues(t, "ping.js loaded from server", result)
+	})
+}
+
 func TestLifecycleReloadLoad(t *testing.T) {
 	t.Parallel()