Integration tests fix

pkg/test/util/ngrok.go

@@ -3,8 +3,10 @@ package util
 import (
 	"context"
 	"fmt"
+	"net/http"
 	"net/url"
 	"os"
+	"time"
 
 	"golang.ngrok.com/ngrok/v2"
 )
@@ -45,6 +47,12 @@ func SetupForwarding(parentCtx context.Context, to string) (*Forwarding, error)
 		return nil, fmt.Errorf("ngrok.Forward failed: %w", err)
 	}
 
+	if err := waitForEdge(ctx, fwd.URL().String()+"/sse"); err != nil {
+		cancel()
+		<-fwd.Done()
+		return nil, fmt.Errorf("ngrok edge readiness probe failed: %w", err)
+	}
+
 	return &Forwarding{
 		URL: fwd.URL(),
 		Shutdown: func() {
@@ -53,3 +61,33 @@ func SetupForwarding(parentCtx context.Context, to string) (*Forwarding, error)
 		},
 	}, nil
 }
+
+// agent.Forward() returns before the ngrok edge is routable, causing 404 on the first request.
+func waitForEdge(ctx context.Context, probeURL string) error {
+	deadline := time.Now().Add(10 * time.Second)
+	backoff := 100 * time.Millisecond
+
+	for time.Now().Before(deadline) {
+		req, err := http.NewRequestWithContext(ctx, http.MethodGet, probeURL, nil)
+		if err != nil {
+			return err
+		}
+		resp, err := http.DefaultClient.Do(req)
+		if err == nil {
+			resp.Body.Close()
+			if resp.StatusCode != http.StatusNotFound {
+				return nil
+			}
+		}
+
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-time.After(backoff):
+		}
+		if backoff < time.Second {
+			backoff *= 2
+		}
+	}
+	return fmt.Errorf("ngrok endpoint %s kept returning 404", probeURL)
+}