github.com/iotexproject/iotex-core@v1.14.1-rc1/api/serverV2_test.go (about)

     1  package api
     2  
     3  import (
     4  	"context"
     5  	"net/http"
     6  	"net/http/httptest"
     7  	"strings"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/golang/mock/gomock"
    12  	"github.com/gorilla/websocket"
    13  	"github.com/pkg/errors"
    14  	"github.com/stretchr/testify/require"
    15  	"golang.org/x/time/rate"
    16  
    17  	"github.com/iotexproject/iotex-core/test/mock/mock_apicoreservice"
    18  	"github.com/iotexproject/iotex-core/testutil"
    19  )
    20  
    21  func TestServerV2(t *testing.T) {
    22  	require := require.New(t)
    23  	ctrl := gomock.NewController(t)
    24  	defer ctrl.Finish()
    25  	core := mock_apicoreservice.NewMockCoreService(ctrl)
    26  	web3Handler := NewWeb3Handler(core, "", _defaultBatchRequestLimit)
    27  	svr := &ServerV2{
    28  		core:         core,
    29  		grpcServer:   NewGRPCServer(core, testutil.RandomPort()),
    30  		httpSvr:      NewHTTPServer("", testutil.RandomPort(), newHTTPHandler(web3Handler)),
    31  		websocketSvr: NewHTTPServer("", testutil.RandomPort(), NewWebsocketHandler(web3Handler, nil)),
    32  	}
    33  	ctx := context.Background()
    34  
    35  	t.Run("start-stop succeed", func(t *testing.T) {
    36  		core.EXPECT().Start(gomock.Any()).Return(nil).Times(1)
    37  		err := svr.Start(ctx)
    38  		require.NoError(err)
    39  
    40  		core.EXPECT().Stop(gomock.Any()).Return(nil).Times(1)
    41  		err = testutil.WaitUntil(100*time.Millisecond, 3*time.Second, func() (bool, error) {
    42  			err = svr.Stop(ctx)
    43  			return err == nil, err
    44  		})
    45  		require.NoError(err)
    46  	})
    47  
    48  	t.Run("start failed", func(t *testing.T) {
    49  		expectErr := errors.New("failed to add chainListener")
    50  		core.EXPECT().Start(gomock.Any()).Return(expectErr).Times(1)
    51  		err := svr.Start(ctx)
    52  		require.Contains(err.Error(), expectErr.Error())
    53  	})
    54  
    55  	t.Run("stop failed", func(t *testing.T) {
    56  		expectErr := errors.New("failed to shutdown api tracer")
    57  		core.EXPECT().Stop(gomock.Any()).Return(expectErr).Times(1)
    58  		err := svr.Stop(ctx)
    59  		require.Contains(err.Error(), expectErr.Error())
    60  	})
    61  
    62  	t.Run("websocket rate limit", func(t *testing.T) {
    63  		// set the limiter to 1 request per second
    64  		limiter := rate.NewLimiter(1, 1)
    65  		echo := func(w http.ResponseWriter, r *http.Request) {
    66  			c, err := upgrader.Upgrade(w, r, nil)
    67  			if err != nil {
    68  				return
    69  			}
    70  			defer c.Close()
    71  			for {
    72  				if err := limiter.Wait(ctx); err != nil {
    73  					return
    74  				}
    75  				mt, message, err := c.ReadMessage()
    76  				if err != nil {
    77  					break
    78  				}
    79  				err = c.WriteMessage(mt, message)
    80  				if err != nil {
    81  					break
    82  				}
    83  			}
    84  		}
    85  		s := httptest.NewServer(http.HandlerFunc(echo))
    86  		defer s.Close()
    87  
    88  		u := "ws" + strings.TrimPrefix(s.URL, "http")
    89  		c, _, err := websocket.DefaultDialer.Dial(u, nil)
    90  		require.NoError(err)
    91  		defer c.Close()
    92  		i := 0
    93  		timeout := time.After(3 * time.Second)
    94  	LOOP:
    95  		for {
    96  			select {
    97  			case <-timeout:
    98  				break LOOP
    99  			default:
   100  				err := c.WriteMessage(websocket.TextMessage, []byte{0})
   101  				require.NoError(err)
   102  				_, _, err = c.ReadMessage()
   103  				require.NoError(err)
   104  				i++
   105  			}
   106  		}
   107  		require.Equal(4, i)
   108  	})
   109  }