github.com/mattermosttest/mattermost-server/v5@v5.0.0-20200917143240-9dfa12e121f9/model/websocket_client_test.go (about)

     1  // Copyright (c) 2015-present Mattermost, Inc. All Rights Reserved.
     2  // See LICENSE.txt for license information.
     3  
     4  package model
     5  
     6  import (
     7  	"net/http"
     8  	"net/http/httptest"
     9  	"strings"
    10  	"testing"
    11  	"time"
    12  
    13  	"github.com/gorilla/websocket"
    14  	"github.com/stretchr/testify/assert"
    15  	"github.com/stretchr/testify/require"
    16  )
    17  
    18  func dummyWebsocketHandler(t *testing.T) http.HandlerFunc {
    19  	return func(w http.ResponseWriter, req *http.Request) {
    20  		upgrader := &websocket.Upgrader{
    21  			ReadBufferSize:  1024,
    22  			WriteBufferSize: 1024,
    23  		}
    24  		conn, err := upgrader.Upgrade(w, req, nil)
    25  		require.Nil(t, err)
    26  		var buf []byte
    27  		for {
    28  			_, buf, err = conn.ReadMessage()
    29  			if err != nil {
    30  				break
    31  			}
    32  			t.Logf("%s\n", buf)
    33  			err = conn.WriteMessage(websocket.PingMessage, []byte("ping"))
    34  			if err != nil {
    35  				break
    36  			}
    37  		}
    38  	}
    39  }
    40  
    41  // TestWebSocketRace needs to be run with -race to verify that
    42  // there is no race.
    43  func TestWebSocketRace(t *testing.T) {
    44  	s := httptest.NewServer(dummyWebsocketHandler(t))
    45  	defer s.Close()
    46  
    47  	url := strings.Replace(s.URL, "http://", "ws://", 1)
    48  	cli, err := NewWebSocketClient4(url, "authToken")
    49  	require.Nil(t, err)
    50  
    51  	cli.Listen()
    52  
    53  	for i := 0; i < 10; i++ {
    54  		time.Sleep(500 * time.Millisecond)
    55  		cli.UserTyping("channel", "parentId")
    56  	}
    57  }
    58  
    59  func TestWebSocketClose(t *testing.T) {
    60  	// This fails in SuddenClose because we check for closing the writeChan
    61  	// only after waiting the closure of Event and Response channels.
    62  	// Therefore, it is still racy and can panic. There is no use chasing this
    63  	// again because it will be completely overhauled in v6.
    64  	t.Skip("Skipping the test. Will be changed in v6.")
    65  	s := httptest.NewServer(dummyWebsocketHandler(t))
    66  	defer s.Close()
    67  
    68  	url := strings.Replace(s.URL, "http://", "ws://", 1)
    69  
    70  	// Check whether the Event and Response channels
    71  	// have been closed or not.
    72  	waitClose := func(doneChan chan struct{}) int {
    73  		numClosed := 0
    74  		timeout := time.After(300 * time.Millisecond)
    75  		for {
    76  			select {
    77  			case <-doneChan:
    78  				numClosed++
    79  				if numClosed == 2 {
    80  					return numClosed
    81  				}
    82  			case <-timeout:
    83  				require.Fail(t, "timed out waiting for channels to be closed")
    84  				return numClosed
    85  			}
    86  		}
    87  	}
    88  
    89  	checkWriteChan := func(writeChan chan writeMessage) {
    90  		defer func() {
    91  			if x := recover(); x == nil {
    92  				require.Fail(t, "should have panicked due to closing a closed channel")
    93  			}
    94  		}()
    95  		close(writeChan)
    96  	}
    97  
    98  	waitForResponses := func(doneChan chan struct{}, cli *WebSocketClient) {
    99  		go func() {
   100  			for range cli.EventChannel {
   101  			}
   102  			doneChan <- struct{}{}
   103  		}()
   104  		go func() {
   105  			for range cli.ResponseChannel {
   106  			}
   107  			doneChan <- struct{}{}
   108  		}()
   109  	}
   110  
   111  	t.Run("SuddenClose", func(t *testing.T) {
   112  		cli, err := NewWebSocketClient4(url, "authToken")
   113  		require.Nil(t, err)
   114  
   115  		cli.Listen()
   116  
   117  		doneChan := make(chan struct{}, 2)
   118  		waitForResponses(doneChan, cli)
   119  
   120  		cli.UserTyping("channelId", "parentId")
   121  		cli.Conn.Close()
   122  
   123  		numClosed := waitClose(doneChan)
   124  		assert.Equal(t, 2, numClosed, "unexpected number of channels closed")
   125  
   126  		// Check whether the write channel was closed or not.
   127  		checkWriteChan(cli.writeChan)
   128  
   129  		require.NotNil(t, cli.ListenError, "non-nil listen error")
   130  		assert.Equal(t, "model.websocket_client.connect_fail.app_error", cli.ListenError.Id, "unexpected error id")
   131  	})
   132  
   133  	t.Run("ExplicitClose", func(t *testing.T) {
   134  		cli, err := NewWebSocketClient4(url, "authToken")
   135  		require.Nil(t, err)
   136  
   137  		cli.Listen()
   138  
   139  		doneChan := make(chan struct{}, 2)
   140  		waitForResponses(doneChan, cli)
   141  
   142  		cli.UserTyping("channelId", "parentId")
   143  		cli.Close()
   144  
   145  		numClosed := waitClose(doneChan)
   146  		assert.Equal(t, 2, numClosed, "unexpected number of channels closed")
   147  
   148  		// Check whether the write channel was closed or not.
   149  		checkWriteChan(cli.writeChan)
   150  	})
   151  }