diff --git a/server_publish_test.go b/server_publish_test.go index 91ba3c35..ad973aba 100644 --- a/server_publish_test.go +++ b/server_publish_test.go @@ -993,10 +993,14 @@ func TestServerPublishErrorTimeout(t *testing.T) { "tls", } { t.Run(proto, func(t *testing.T) { + connClosed := make(chan struct{}) sessionClosed := make(chan struct{}) s := &Server{ Handler: &testServerHandler{ + onConnClose: func(sc *ServerConn, err error) { + close(connClosed) + }, onSessionClose: func(ss *ServerSession, err error) { close(sessionClosed) }, @@ -1126,6 +1130,10 @@ func TestServerPublishErrorTimeout(t *testing.T) { require.Equal(t, base.StatusOK, res.StatusCode) <-sessionClosed + + if proto == "tcp" { + <-connClosed + } }) } }