diff --git a/app/web_conn.go b/app/web_conn.go
index 198e0edbd0..8c4b0b706b 100644
--- a/app/web_conn.go
+++ b/app/web_conn.go
@@ -307,8 +307,6 @@ func (wc *WebConn) SetSession(v *model.Session) {
 // Pump starts the WebConn instance. After this, the websocket
 // is ready to send/receive messages.
 func (wc *WebConn) Pump() {
-	defer wc.App.Srv().userService.ReturnSessionToPool(wc.GetSession())
-
 	var wg sync.WaitGroup
 	wg.Add(1)
 	go func() {
diff --git a/app/web_hub.go b/app/web_hub.go
index 5591d06073..b73ed86f38 100644
--- a/app/web_hub.go
+++ b/app/web_hub.go
@@ -630,6 +630,8 @@ func (i *hubConnectionIndex) Add(wc *WebConn) {
 }
 
 func (i *hubConnectionIndex) Remove(wc *WebConn) {
+	wc.App.Srv().userService.ReturnSessionToPool(wc.GetSession())
+
 	userConnIndex, ok := i.byConnection[wc]
 	if !ok {
 		return
diff --git a/app/web_hub_test.go b/app/web_hub_test.go
index e53a2b51d0..470553791c 100644
--- a/app/web_hub_test.go
+++ b/app/web_hub_test.go
@@ -229,6 +229,7 @@ func TestHubConnIndex(t *testing.T) {
 		UserId: model.NewId(),
 	}
 	wc1.SetConnectionID(model.NewId())
+	wc1.SetSession(&model.Session{})
 
 	// User2
 	wc2 := &WebConn{
@@ -236,16 +237,21 @@ func TestHubConnIndex(t *testing.T) {
 		UserId: model.NewId(),
 	}
 	wc2.SetConnectionID(model.NewId())
+	wc2.SetSession(&model.Session{})
+
 	wc3 := &WebConn{
 		App:    th.App,
 		UserId: wc2.UserId,
 	}
 	wc3.SetConnectionID(model.NewId())
+	wc3.SetSession(&model.Session{})
+
 	wc4 := &WebConn{
 		App:    th.App,
 		UserId: wc2.UserId,
 	}
 	wc4.SetConnectionID(model.NewId())
+	wc4.SetSession(&model.Session{})
 
 	connIndex.Add(wc1)
 	connIndex.Add(wc2)
@@ -309,6 +315,7 @@ func TestHubConnIndexByConnectionId(t *testing.T) {
 		UserId: model.NewId(),
 	}
 	wc1.SetConnectionID(wc1ID)
+	wc1.SetSession(&model.Session{})
 
 	// User2
 	wc2ID := model.NewId()
@@ -317,6 +324,7 @@ func TestHubConnIndexByConnectionId(t *testing.T) {
 		UserId: model.NewId(),
 	}
 	wc2.SetConnectionID(wc2ID)
+	wc2.SetSession(&model.Session{})
 
 	wc3ID := model.NewId()
 	wc3 := &WebConn{
@@ -324,6 +332,7 @@ func TestHubConnIndexByConnectionId(t *testing.T) {
 		UserId: wc2.UserId,
 	}
 	wc3.SetConnectionID(wc3ID)
+	wc3.SetSession(&model.Session{})
 
 	t.Run("no connections", func(t *testing.T) {
 		assert.False(t, connIndex.Has(wc1))
@@ -353,26 +362,36 @@ func TestHubConnIndexByConnectionId(t *testing.T) {
 }
 
 func TestHubConnIndexInactive(t *testing.T) {
+	th := Setup(t)
+	defer th.TearDown()
+
 	connIndex := newHubConnectionIndex(2 * time.Second)
 
 	// User1
 	wc1 := &WebConn{
+		App:    th.App,
 		UserId: model.NewId(),
 		active: true,
 	}
 	wc1.SetConnectionID("conn1")
+	wc1.SetSession(&model.Session{})
 
 	// User2
 	wc2 := &WebConn{
+		App:    th.App,
 		UserId: model.NewId(),
 		active: true,
 	}
 	wc2.SetConnectionID("conn2")
+	wc2.SetSession(&model.Session{})
+
 	wc3 := &WebConn{
+		App:    th.App,
 		UserId: wc2.UserId,
 		active: false,
 	}
 	wc3.SetConnectionID("conn3")
+	wc3.SetSession(&model.Session{})
 
 	connIndex.Add(wc1)
 	connIndex.Add(wc2)