From 6a9908bda7aa49378cfa170d75fdf0db496fdeff Mon Sep 17 00:00:00 2001
From: "google-labs-jules[bot]"
 <161369871+google-labs-jules[bot]@users.noreply.github.com>
Date: Sun, 21 Dec 2025 06:32:24 +0000
Subject: [PATCH] Fix TestNewShellHandler to use raw byte streams

Updated the websocket test to expect raw byte messages instead of JSON
envelopes with base64 encoded data, aligning the test with the current
implementation of the websocket handler. This resolves test failures
where the test expected `{"type":"data",...}` but received raw text.
---
 tavern/internal/http/stream/websocket_test.go | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/tavern/internal/http/stream/websocket_test.go b/tavern/internal/http/stream/websocket_test.go
index 01006d32c..f3dce0b81 100644
--- a/tavern/internal/http/stream/websocket_test.go
+++ b/tavern/internal/http/stream/websocket_test.go
@@ -104,12 +104,11 @@ func TestNewShellHandler(t *testing.T) {
 	_, p, err := ws.ReadMessage()
 	assert.NoError(t, err)
 
-	expectedJSON := `{"type":"data","data":"aGVsbG8gZnJvbSBzZXJ2ZXI="}` // base64 "hello from server"
-	assert.Equal(t, []byte(expectedJSON), p)
+	assert.Equal(t, testMessage, p)
 
 	// Test reading from the websocket (shell -> server)
-	// Client sends JSON
-	readMessage := []byte(`{"type":"data","data":"aGVsbG8gZnJvbSBzaGVsbA=="}`) // base64 "hello from shell"
+	// Client sends raw bytes
+	readMessage := []byte("hello from shell")
 	err = ws.WriteMessage(websocket.TextMessage, readMessage)
 	require.NoError(t, err)
 
@@ -118,7 +117,7 @@ func TestNewShellHandler(t *testing.T) {
 	require.NoError(t, err, "timed out waiting for message from websocket")
 
 	// The body sent to pubsub should be the raw bytes
-	assert.Equal(t, []byte("hello from shell"), msg.Body)
+	assert.Equal(t, readMessage, msg.Body)
 	assert.Equal(t, "data", msg.Metadata[stream.MetadataMsgKind])
 	msg.Ack()
 }