github.com/vnforks/kid@v5.11.1+incompatible/app/web_hub_test.go (about)

     1  package app
     2  
     3  import (
     4  	"net"
     5  	"net/http"
     6  	"net/http/httptest"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/gorilla/websocket"
    11  	goi18n "github.com/nicksnyder/go-i18n/i18n"
    12  	"github.com/stretchr/testify/require"
    13  
    14  	"github.com/mattermost/mattermost-server/model"
    15  )
    16  
    17  func dummyWebsocketHandler(t *testing.T) http.HandlerFunc {
    18  	return func(w http.ResponseWriter, req *http.Request) {
    19  		upgrader := &websocket.Upgrader{
    20  			ReadBufferSize:  1024,
    21  			WriteBufferSize: 1024,
    22  		}
    23  		conn, err := upgrader.Upgrade(w, req, nil)
    24  		for err == nil {
    25  			_, _, err = conn.ReadMessage()
    26  		}
    27  		if _, ok := err.(*websocket.CloseError); !ok {
    28  			require.NoError(t, err)
    29  		}
    30  	}
    31  }
    32  
    33  func registerDummyWebConn(t *testing.T, a *App, addr net.Addr, userId string) *WebConn {
    34  	session, appErr := a.CreateSession(&model.Session{
    35  		UserId: userId,
    36  	})
    37  	require.Nil(t, appErr)
    38  
    39  	d := websocket.Dialer{}
    40  	c, _, err := d.Dial("ws://"+addr.String()+"/ws", nil)
    41  	require.NoError(t, err)
    42  
    43  	wc := a.NewWebConn(c, *session, goi18n.IdentityTfunc(), "en")
    44  	a.HubRegister(wc)
    45  	go wc.Pump()
    46  	return wc
    47  }
    48  
    49  func TestHubStopWithMultipleConnections(t *testing.T) {
    50  	th := Setup(t).InitBasic()
    51  	defer th.TearDown()
    52  
    53  	s := httptest.NewServer(http.HandlerFunc(dummyWebsocketHandler(t)))
    54  	defer s.Close()
    55  
    56  	th.App.HubStart()
    57  	wc1 := registerDummyWebConn(t, th.App, s.Listener.Addr(), th.BasicUser.Id)
    58  	wc2 := registerDummyWebConn(t, th.App, s.Listener.Addr(), th.BasicUser.Id)
    59  	wc3 := registerDummyWebConn(t, th.App, s.Listener.Addr(), th.BasicUser.Id)
    60  	defer wc1.Close()
    61  	defer wc2.Close()
    62  	defer wc3.Close()
    63  }
    64  
    65  // TestHubStopRaceCondition verifies that attempts to use the hub after it has shutdown does not
    66  // block the caller indefinitely.
    67  func TestHubStopRaceCondition(t *testing.T) {
    68  	th := Setup(t).InitBasic()
    69  	defer th.TearDown()
    70  
    71  	s := httptest.NewServer(http.HandlerFunc(dummyWebsocketHandler(t)))
    72  
    73  	th.App.HubStart()
    74  	wc1 := registerDummyWebConn(t, th.App, s.Listener.Addr(), th.BasicUser.Id)
    75  	defer wc1.Close()
    76  
    77  	hub := th.App.Srv.Hubs[0]
    78  	th.App.HubStop()
    79  	time.Sleep(5 * time.Second)
    80  
    81  	done := make(chan bool)
    82  	go func() {
    83  		wc4 := registerDummyWebConn(t, th.App, s.Listener.Addr(), th.BasicUser.Id)
    84  		wc5 := registerDummyWebConn(t, th.App, s.Listener.Addr(), th.BasicUser.Id)
    85  		hub.Register(wc4)
    86  		hub.Register(wc5)
    87  
    88  		hub.UpdateActivity("userId", "sessionToken", 0)
    89  
    90  		for i := 0; i <= BROADCAST_QUEUE_SIZE; i++ {
    91  			hub.Broadcast(&model.WebSocketEvent{})
    92  		}
    93  
    94  		hub.InvalidateUser("userId")
    95  		hub.Unregister(wc4)
    96  		hub.Unregister(wc5)
    97  		close(done)
    98  	}()
    99  
   100  	select {
   101  	case <-done:
   102  	case <-time.After(15 * time.Second):
   103  		t.Fatalf("hub call did not return within 15 seconds after stop")
   104  	}
   105  }