github.com/ari-anchor/sei-tendermint@v0.0.0-20230519144642-dc826b7b56bb/rpc/jsonrpc/server/ws_handler_test.go (about)

     1  package server
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"net/http"
     7  	"net/http/httptest"
     8  	"testing"
     9  
    10  	"github.com/fortytw2/leaktest"
    11  	"github.com/gorilla/websocket"
    12  	"github.com/stretchr/testify/require"
    13  
    14  	"github.com/ari-anchor/sei-tendermint/libs/log"
    15  	rpctypes "github.com/ari-anchor/sei-tendermint/rpc/jsonrpc/types"
    16  )
    17  
    18  func TestWebsocketManagerHandler(t *testing.T) {
    19  	logger := log.NewNopLogger()
    20  
    21  	s := newWSServer(t, logger)
    22  	defer s.Close()
    23  
    24  	t.Cleanup(leaktest.Check(t))
    25  
    26  	// check upgrader works
    27  	d := websocket.Dialer{}
    28  	c, dialResp, err := d.Dial("ws://"+s.Listener.Addr().String()+"/websocket", nil)
    29  	require.NoError(t, err)
    30  
    31  	if got, want := dialResp.StatusCode, http.StatusSwitchingProtocols; got != want {
    32  		t.Errorf("dialResp.StatusCode = %q, want %q", got, want)
    33  	}
    34  
    35  	// check basic functionality works
    36  	req := rpctypes.NewRequest(1001)
    37  	require.NoError(t, req.SetMethodAndParams("c", map[string]interface{}{"s": "a", "i": 10}))
    38  	require.NoError(t, c.WriteJSON(req))
    39  
    40  	var resp rpctypes.RPCResponse
    41  	err = c.ReadJSON(&resp)
    42  	require.NoError(t, err)
    43  	require.Nil(t, resp.Error)
    44  	dialResp.Body.Close()
    45  }
    46  
    47  func newWSServer(t *testing.T, logger log.Logger) *httptest.Server {
    48  	type args struct {
    49  		S string      `json:"s"`
    50  		I json.Number `json:"i"`
    51  	}
    52  	funcMap := map[string]*RPCFunc{
    53  		"c": NewWSRPCFunc(func(context.Context, *args) (string, error) { return "foo", nil }),
    54  	}
    55  	wm := NewWebsocketManager(logger, funcMap)
    56  
    57  	mux := http.NewServeMux()
    58  	mux.HandleFunc("/websocket", wm.WebsocketHandler)
    59  
    60  	srv := httptest.NewServer(mux)
    61  
    62  	t.Cleanup(srv.Close)
    63  
    64  	return srv
    65  }