chore: improve test

xmidt-org · Apr 3, 2024 · ba86d4a · ba86d4a
1 parent 239e04c
commit ba86d4a
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/internal/websocket/e2e_test.go b/internal/websocket/e2e_test.go
@@ -5,6 +5,7 @@ package websocket_test
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"net/http"
 	"net/http/httptest"
@@ -59,7 +60,9 @@ func (ts *EndToEndTestSuite) TestEndToEnd() {
 
 				mt, got, err := c.Read(ctx)
 				// server will halt until the websocket closes resulting in a EOF
-				if ts.finished {
+				var closeErr websocket.CloseError
+				if ts.finished && errors.As(err, &closeErr) {
+					assert.Equal(closeErr.Code, websocket.StatusNormalClosure)
 					return
 				}