fix a race in tests

Signed-off-by: Dmitri Dolguikh <dmitri.external@netbird.io>

util/wsproxy/server/proxy.go

@@ -2,6 +2,7 @@ package server
 
 import (
 	"net/http"
+	"sync/atomic"
 	"time"
 
 	"github.com/coder/websocket"
@@ -59,7 +60,7 @@ func (p *Proxy) Handler() http.Handler {
 type proxyHandler struct {
 	metrics            MetricsRecorder
 	handler            http.Handler
-	conn               *wsConnAdapter
+	conn               atomic.Pointer[wsConnAdapter]
 	headersReadTimeout time.Duration
 }
 
@@ -91,7 +92,7 @@ func (ph *proxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		_ = serverConn.Close()
 	}()
 
-	ph.conn = serverConn // used in tests only
+	ph.conn.Store(serverConn) // used in tests only
 
 	log.Debugf("WebSocket proxy established: %s -> gRPC handler", r.RemoteAddr)
 

util/wsproxy/server/ws_conn_adapter_test.go

@@ -90,7 +90,7 @@ func TestAdapterHandlingConnectionClosures(t *testing.T) {
 			case 0:
 				clientconn.Close(websocket.StatusNormalClosure, "")
 			case 1:
-				handler.conn.Close()
+				handler.conn.Load().Close()
 			case 2:
 				cancel()
 			case 3:
@@ -99,7 +99,7 @@ func TestAdapterHandlingConnectionClosures(t *testing.T) {
 			}
 
 			assert.EventuallyWithT(t, func(c *assert.CollectT) {
-				assert.True(c, handler.conn.IsClosed())
+				assert.True(c, handler.conn.Load().IsClosed())
 			}, 3*time.Second, 100*time.Millisecond)
 		})
 	}
@@ -154,7 +154,7 @@ func TestAdapterHandlingHttpConnection_NoHeadersSent(t *testing.T) {
 	assert.Error(t, err)
 
 	assert.EventuallyWithT(t, func(c *assert.CollectT) {
-		assert.True(c, handler.conn.IsClosed())
+		assert.True(c, handler.conn.Load().IsClosed())
 	}, 3*time.Second, 100*time.Millisecond)
 }