k8s-operator: fix test flake (#16680)

This occasionally panics waiting on a nil ctx, but was missed in the
previous PR because it's quite a rare flake as it needs to progress to a
specific point in the parser.

Updates #16678

Change-Id: Ifd36dfc915b153aede36b8ee39eff83750031f95

Signed-off-by: Tom Proctor <tomhjp@users.noreply.github.com>

k8s-operator/sessionrecording/ws/conn_test.go

@@ -9,6 +9,7 @@ import (
 	"context"
 	"fmt"
 	"reflect"
+	"runtime/debug"
 	"testing"
 	"time"
 
@@ -284,19 +285,28 @@ func Test_conn_WriteRand(t *testing.T) {
 	sr := &fakes.TestSessionRecorder{}
 	rec := tsrecorder.New(sr, cl, cl.Now(), true, zl.Sugar())
 	for i := range 100 {
-		tc := &fakes.TestConn{}
+		t.Run(fmt.Sprintf("test_%d", i), func(t *testing.T) {
-		c := &conn{
+			tc := &fakes.TestConn{}
-			Conn: tc,
+			c := &conn{
-			log:  zl.Sugar(),
+				Conn: tc,
-			rec:  rec,
+				log:  zl.Sugar(),
-		}
+				rec:  rec,
-		bb := fakes.RandomBytes(t)
+
-		for j, input := range bb {
+				ctx:                   context.Background(), // ctx must be non-nil.
-			f := func() {
+				initialCastHeaderSent: make(chan struct{}),
-				c.Write(input)
 			}
-			testPanic(t, f, fmt.Sprintf("[%d %d] Write: panic parsing input of length %d first bytes %b current write message %+#v", i, j, len(input), firstBytes(input), c.currentWriteMsg))
+			// Never block for random data.
-		}
+			c.writeCastHeaderOnce.Do(func() {
+				close(c.initialCastHeaderSent)
+			})
+			bb := fakes.RandomBytes(t)
+			for j, input := range bb {
+				f := func() {
+					c.Write(input)
+				}
+				testPanic(t, f, fmt.Sprintf("[%d %d] Write: panic parsing input of length %d first bytes %b current write message %+#v", i, j, len(input), firstBytes(input), c.currentWriteMsg))
+			}
+		})
 	}
 }
 
@@ -304,7 +314,7 @@ func testPanic(t *testing.T, f func(), msg string) {
 	t.Helper()
 	defer func() {
 		if r := recover(); r != nil {
-			t.Fatal(msg, r)
+			t.Fatal(msg, r, string(debug.Stack()))
 		}
 	}()
 	f()