From f45d28d4a43f70e1aff7fc9091b780c5a2a61e64 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Maximilian=20Pa=C3=9F?=
 <22845248+mpass99@users.noreply.github.com>
Date: Fri, 10 Dec 2021 12:43:52 +0100
Subject: [PATCH] Fix flaky TestOutputToStdout test

---
 tests/e2e/websocket_test.go | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/tests/e2e/websocket_test.go b/tests/e2e/websocket_test.go
index f6bc895..ebd1735 100644
--- a/tests/e2e/websocket_test.go
+++ b/tests/e2e/websocket_test.go
@@ -63,9 +63,12 @@ func (s *E2ETestSuite) TestOutputToStdout() {
 	s.Require().Error(err)
 	s.Equal(&websocket.CloseError{Code: websocket.CloseNormalClosure}, err)
 
-	s.Require().Equal(&dto.WebSocketMessage{Type: dto.WebSocketMetaStart}, messages[0])
-	s.Require().Equal(&dto.WebSocketMessage{Type: dto.WebSocketOutputStdout, Data: "Hello World\r\n"}, messages[1])
-	s.Require().Equal(&dto.WebSocketMessage{Type: dto.WebSocketExit}, messages[2])
+	controlMessages := helpers.WebSocketControlMessages(messages)
+	s.Require().Equal(&dto.WebSocketMessage{Type: dto.WebSocketMetaStart}, controlMessages[0])
+	s.Require().Equal(&dto.WebSocketMessage{Type: dto.WebSocketExit}, controlMessages[1])
+
+	stdout, _, _ := helpers.WebSocketOutputMessages(messages)
+	s.Require().Equal("Hello World\r\n", stdout)
 }
 
 func (s *E2ETestSuite) TestOutputToStderr() {