From 89a99a08bfd7dc54ed9876fce72af32b56c10024 Mon Sep 17 00:00:00 2001 From: Toby Date: Sat, 23 Mar 2024 11:17:51 -0700 Subject: [PATCH] fix: flaky tests caused by occasionally closing channel multiple times --- core/internal/integration_tests/close_test.go | 7 +++++-- core/internal/integration_tests/trafficlogger_test.go | 11 +++++++---- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/core/internal/integration_tests/close_test.go b/core/internal/integration_tests/close_test.go index 4160b3c..ab6d016 100644 --- a/core/internal/integration_tests/close_test.go +++ b/core/internal/integration_tests/close_test.go @@ -2,6 +2,7 @@ package integration_tests import ( "io" + "sync" "testing" "time" @@ -48,13 +49,14 @@ func TestClientServerTCPClose(t *testing.T) { // Server outbound connection should write the same thing, then close. sobConn := mocks.NewMockConn(t) sobConnCh := make(chan struct{}) // For close signal only + sobConnChCloseFunc := sync.OnceFunc(func() { close(sobConnCh) }) sobConn.EXPECT().Read(mock.Anything).RunAndReturn(func(bs []byte) (int, error) { <-sobConnCh return 0, io.EOF }) sobConn.EXPECT().Write([]byte("happy")).Return(5, nil) sobConn.EXPECT().Close().RunAndReturn(func() error { - close(sobConnCh) + sobConnChCloseFunc() return nil }) serverOb.EXPECT().TCP(addr).Return(sobConn, nil).Once() @@ -133,6 +135,7 @@ func TestClientServerUDPIdleTimeout(t *testing.T) { // to trigger the server's UDP idle timeout. sobConn := mocks.NewMockUDPConn(t) sobConnCh := make(chan []byte, 1) + sobConnChCloseFunc := sync.OnceFunc(func() { close(sobConnCh) }) sobConn.EXPECT().ReadFrom(mock.Anything).RunAndReturn(func(bs []byte) (int, string, error) { d := <-sobConnCh if d == nil { @@ -167,7 +170,7 @@ func TestClientServerUDPIdleTimeout(t *testing.T) { } // Now we wait for 3 seconds, the server should close the UDP session. sobConn.EXPECT().Close().RunAndReturn(func() error { - close(sobConnCh) + sobConnChCloseFunc() return nil }) eventLogger.EXPECT().UDPError(mock.Anything, mock.Anything, uint32(1), nil).Once() diff --git a/core/internal/integration_tests/trafficlogger_test.go b/core/internal/integration_tests/trafficlogger_test.go index ff1d66e..abbea8c 100644 --- a/core/internal/integration_tests/trafficlogger_test.go +++ b/core/internal/integration_tests/trafficlogger_test.go @@ -2,6 +2,7 @@ package integration_tests import ( "io" + "sync" "testing" "time" @@ -46,6 +47,7 @@ func TestClientServerTrafficLoggerTCP(t *testing.T) { sobConn := mocks.NewMockConn(t) sobConnCh := make(chan []byte, 1) + sobConnChCloseFunc := sync.OnceFunc(func() { close(sobConnCh) }) sobConn.EXPECT().Read(mock.Anything).RunAndReturn(func(bs []byte) (int, error) { b := <-sobConnCh if b == nil { @@ -55,9 +57,9 @@ func TestClientServerTrafficLoggerTCP(t *testing.T) { } }) sobConn.EXPECT().Close().RunAndReturn(func() error { - close(sobConnCh) + sobConnChCloseFunc() return nil - }).Once() + }) serverOb.EXPECT().TCP(addr).Return(sobConn, nil).Once() conn, err := c.TCP(addr) @@ -125,6 +127,7 @@ func TestClientServerTrafficLoggerUDP(t *testing.T) { sobConn := mocks.NewMockUDPConn(t) sobConnCh := make(chan []byte, 1) + sobConnChCloseFunc := sync.OnceFunc(func() { close(sobConnCh) }) sobConn.EXPECT().ReadFrom(mock.Anything).RunAndReturn(func(bs []byte) (int, string, error) { b := <-sobConnCh if b == nil { @@ -134,9 +137,9 @@ func TestClientServerTrafficLoggerUDP(t *testing.T) { } }) sobConn.EXPECT().Close().RunAndReturn(func() error { - close(sobConnCh) + sobConnChCloseFunc() return nil - }).Once() + }) serverOb.EXPECT().UDP(addr).Return(sobConn, nil).Once() conn, err := c.UDP()