From 75701139dd4e8d82a52584bbf83742823a8be744 Mon Sep 17 00:00:00 2001
From: schmidtw <weston_schmidt@alumni.purdue.edu>
Date: Sat, 30 Dec 2023 14:12:21 -0800
Subject: [PATCH] Fix the race condition in the end to end test.

---
 internal/websocket/e2e_test.go | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/internal/websocket/e2e_test.go b/internal/websocket/e2e_test.go
index a3ee89c..0606561 100644
--- a/internal/websocket/e2e_test.go
+++ b/internal/websocket/e2e_test.go
@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/http/httptest"
+	"sync"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -240,8 +241,12 @@ func TestEndToEndConnectionIssues(t *testing.T) {
 
 	var msgCnt, connectCnt, disconnectCnt atomic.Int64
 
+	var mutex sync.Mutex
+
 	got, err := ws.New(
 		ws.FetchURL(func(context.Context) (string, error) {
+			mutex.Lock()
+			defer mutex.Unlock()
 			if s.URL == "" {
 				return "", fmt.Errorf("no url")
 			}
@@ -280,16 +285,20 @@ func TestEndToEndConnectionIssues(t *testing.T) {
 	var started bool
 	for {
 		if connectCnt.Load() >= 3 && !started {
+			mutex.Lock()
 			s.Start()
-			defer s.Close()
+			mutex.Unlock()
+			defer func() {
+				mutex.Lock()
+				s.Close()
+				mutex.Unlock()
+			}()
 			started = true
 		}
 		if disconnectCnt.Load() > 0 {
 			break
 		}
 
-		fmt.Printf("connectCnt: %d, disconnectCnt: %d, msgCnt: %d\n", connectCnt.Load(), disconnectCnt.Load(), msgCnt.Load())
-
 		select {
 		case <-ctx.Done():
 			assert.Fail("timed out waiting for messages")