fix: flaky tests caused by occasionally closing channel multiple times

This commit is contained in:
Toby 2024-03-23 11:17:51 -07:00
parent a037880f88
commit 89a99a08bf
2 changed files with 12 additions and 6 deletions

View file

@ -2,6 +2,7 @@ package integration_tests
import ( import (
"io" "io"
"sync"
"testing" "testing"
"time" "time"
@ -48,13 +49,14 @@ func TestClientServerTCPClose(t *testing.T) {
// Server outbound connection should write the same thing, then close. // Server outbound connection should write the same thing, then close.
sobConn := mocks.NewMockConn(t) sobConn := mocks.NewMockConn(t)
sobConnCh := make(chan struct{}) // For close signal only 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) { sobConn.EXPECT().Read(mock.Anything).RunAndReturn(func(bs []byte) (int, error) {
<-sobConnCh <-sobConnCh
return 0, io.EOF return 0, io.EOF
}) })
sobConn.EXPECT().Write([]byte("happy")).Return(5, nil) sobConn.EXPECT().Write([]byte("happy")).Return(5, nil)
sobConn.EXPECT().Close().RunAndReturn(func() error { sobConn.EXPECT().Close().RunAndReturn(func() error {
close(sobConnCh) sobConnChCloseFunc()
return nil return nil
}) })
serverOb.EXPECT().TCP(addr).Return(sobConn, nil).Once() 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. // to trigger the server's UDP idle timeout.
sobConn := mocks.NewMockUDPConn(t) sobConn := mocks.NewMockUDPConn(t)
sobConnCh := make(chan []byte, 1) sobConnCh := make(chan []byte, 1)
sobConnChCloseFunc := sync.OnceFunc(func() { close(sobConnCh) })
sobConn.EXPECT().ReadFrom(mock.Anything).RunAndReturn(func(bs []byte) (int, string, error) { sobConn.EXPECT().ReadFrom(mock.Anything).RunAndReturn(func(bs []byte) (int, string, error) {
d := <-sobConnCh d := <-sobConnCh
if d == nil { 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. // Now we wait for 3 seconds, the server should close the UDP session.
sobConn.EXPECT().Close().RunAndReturn(func() error { sobConn.EXPECT().Close().RunAndReturn(func() error {
close(sobConnCh) sobConnChCloseFunc()
return nil return nil
}) })
eventLogger.EXPECT().UDPError(mock.Anything, mock.Anything, uint32(1), nil).Once() eventLogger.EXPECT().UDPError(mock.Anything, mock.Anything, uint32(1), nil).Once()

View file

@ -2,6 +2,7 @@ package integration_tests
import ( import (
"io" "io"
"sync"
"testing" "testing"
"time" "time"
@ -46,6 +47,7 @@ func TestClientServerTrafficLoggerTCP(t *testing.T) {
sobConn := mocks.NewMockConn(t) sobConn := mocks.NewMockConn(t)
sobConnCh := make(chan []byte, 1) sobConnCh := make(chan []byte, 1)
sobConnChCloseFunc := sync.OnceFunc(func() { close(sobConnCh) })
sobConn.EXPECT().Read(mock.Anything).RunAndReturn(func(bs []byte) (int, error) { sobConn.EXPECT().Read(mock.Anything).RunAndReturn(func(bs []byte) (int, error) {
b := <-sobConnCh b := <-sobConnCh
if b == nil { if b == nil {
@ -55,9 +57,9 @@ func TestClientServerTrafficLoggerTCP(t *testing.T) {
} }
}) })
sobConn.EXPECT().Close().RunAndReturn(func() error { sobConn.EXPECT().Close().RunAndReturn(func() error {
close(sobConnCh) sobConnChCloseFunc()
return nil return nil
}).Once() })
serverOb.EXPECT().TCP(addr).Return(sobConn, nil).Once() serverOb.EXPECT().TCP(addr).Return(sobConn, nil).Once()
conn, err := c.TCP(addr) conn, err := c.TCP(addr)
@ -125,6 +127,7 @@ func TestClientServerTrafficLoggerUDP(t *testing.T) {
sobConn := mocks.NewMockUDPConn(t) sobConn := mocks.NewMockUDPConn(t)
sobConnCh := make(chan []byte, 1) sobConnCh := make(chan []byte, 1)
sobConnChCloseFunc := sync.OnceFunc(func() { close(sobConnCh) })
sobConn.EXPECT().ReadFrom(mock.Anything).RunAndReturn(func(bs []byte) (int, string, error) { sobConn.EXPECT().ReadFrom(mock.Anything).RunAndReturn(func(bs []byte) (int, string, error) {
b := <-sobConnCh b := <-sobConnCh
if b == nil { if b == nil {
@ -134,9 +137,9 @@ func TestClientServerTrafficLoggerUDP(t *testing.T) {
} }
}) })
sobConn.EXPECT().Close().RunAndReturn(func() error { sobConn.EXPECT().Close().RunAndReturn(func() error {
close(sobConnCh) sobConnChCloseFunc()
return nil return nil
}).Once() })
serverOb.EXPECT().UDP(addr).Return(sobConn, nil).Once() serverOb.EXPECT().UDP(addr).Return(sobConn, nil).Once()
conn, err := c.UDP() conn, err := c.UDP()