From f2af7489ff223bfa14f39b5913dc7b74fc1c5923 Mon Sep 17 00:00:00 2001
From: aler9 <46489434+aler9@users.noreply.github.com>
Date: Tue, 4 May 2021 15:59:51 +0200
Subject: [PATCH] server: test connection timeout

---
 server_publish_test.go | 8 ++++++++
 1 file changed, 8 insertions(+)

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
+			}
 		})
 	}
 }