From 3435a95a761cb5aac45292893e92f975e97848ba Mon Sep 17 00:00:00 2001 From: Peter Broadhurst Date: Fri, 28 Jan 2022 09:14:28 -0500 Subject: [PATCH] Avoid potential double-close of channels due to fast reconnects in test Signed-off-by: Peter Broadhurst --- pkg/wsclient/wsclient_test.go | 5 +++++ pkg/wsclient/wstestserver.go | 5 +++++ 2 files changed, 10 insertions(+) diff --git a/pkg/wsclient/wsclient_test.go b/pkg/wsclient/wsclient_test.go index 3180203dfe..1440f8f92b 100644 --- a/pkg/wsclient/wsclient_test.go +++ b/pkg/wsclient/wsclient_test.go @@ -201,6 +201,7 @@ func TestWSReadLoopSendFailure(t *testing.T) { defer done() wsconn, _, err := websocket.DefaultDialer.Dial(url, nil) + assert.NoError(t, err) wsconn.WriteJSON(map[string]string{"type": "listen", "topic": "topic1"}) assert.NoError(t, err) <-toServer @@ -217,6 +218,10 @@ func TestWSReadLoopSendFailure(t *testing.T) { // Ensure the readLoop exits immediately w.readLoop() + // Try reconnect, should fail here + _, _, err = websocket.DefaultDialer.Dial(url, nil) + assert.Error(t, err) + } func TestWSReconnectFail(t *testing.T) { diff --git a/pkg/wsclient/wstestserver.go b/pkg/wsclient/wstestserver.go index a487c03d9f..aa7ee9149d 100644 --- a/pkg/wsclient/wstestserver.go +++ b/pkg/wsclient/wstestserver.go @@ -36,6 +36,11 @@ func NewTestWSServer(testReq func(req *http.Request)) (toServer, fromServer chan if testReq != nil { testReq(req) } + if connected { + // test server only handles one open connection, as it only has one set of channels + res.WriteHeader(409) + return + } ws, _ := upgrader.Upgrade(res, req, http.Header{}) go func() { defer close(receiveDone)