github.com/polygon-io/client-go@v1.16.4/websocket/polygon_test.go (about)

     1  package polygonws
     2  
     3  import (
     4  	"encoding/json"
     5  	"net/http"
     6  	"net/http/httptest"
     7  	"strings"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/gorilla/websocket"
    12  	"github.com/polygon-io/client-go/websocket/models"
    13  	"github.com/sirupsen/logrus"
    14  	"github.com/stretchr/testify/assert"
    15  )
    16  
    17  func connect(w http.ResponseWriter, r *http.Request) {
    18  	upgrader := websocket.Upgrader{}
    19  	c, err := upgrader.Upgrade(w, r, nil)
    20  	if err != nil {
    21  		return
    22  	}
    23  	defer c.Close()
    24  
    25  	for {
    26  		mt, msg, err := c.ReadMessage()
    27  		if err != nil {
    28  			return
    29  		}
    30  
    31  		var cm models.ControlMessage
    32  		_ = json.Unmarshal(msg, &cm)
    33  		if cm.Action == "auth" && cm.Params == "good" {
    34  			res := []models.ControlMessage{{EventType: models.EventType{EventType: "status"}, Status: "auth_success"}}
    35  			data, _ := json.Marshal(res)
    36  			err = c.WriteMessage(mt, data)
    37  		} else {
    38  			res := []models.ControlMessage{{EventType: models.EventType{EventType: "status"}, Status: "auth_failed"}}
    39  			data, _ := json.Marshal(res)
    40  			err = c.WriteMessage(mt, data)
    41  		}
    42  		if err != nil {
    43  			return
    44  		}
    45  	}
    46  }
    47  
    48  func TestNew(t *testing.T) {
    49  	// successful creation
    50  	c, err := New(Config{
    51  		APIKey: "test",
    52  		Feed:   PolyFeed,
    53  		Market: Options,
    54  	})
    55  	assert.NotNil(t, c)
    56  	assert.Nil(t, err)
    57  	assert.Equal(t, "wss://polyfeed.polygon.io/options", c.url)
    58  	assert.Equal(t, &nopLogger{}, c.log)
    59  
    60  	// empty config
    61  	c, err = New(Config{})
    62  	assert.Nil(t, c)
    63  	assert.NotNil(t, err)
    64  }
    65  
    66  func TestConnectAuthSuccess(t *testing.T) {
    67  	s := httptest.NewServer(http.HandlerFunc(connect))
    68  	defer s.Close()
    69  
    70  	log := logrus.New()
    71  	log.SetLevel(logrus.DebugLevel)
    72  	u := "ws" + strings.TrimPrefix(s.URL, "http")
    73  	var retries uint64 = 0
    74  	c, err := New(Config{
    75  		APIKey:     "good",
    76  		Feed:       Feed(u),
    77  		Market:     Market(""),
    78  		Log:        log,
    79  		MaxRetries: &retries,
    80  	})
    81  	assert.NotNil(t, c)
    82  	assert.Nil(t, err)
    83  
    84  	defer func() {
    85  		time.Sleep(100 * time.Millisecond)
    86  		c.Close()
    87  	}()
    88  
    89  	// closing before connecting shouldn't do anthing
    90  	c.Close()
    91  
    92  	// connect successfully
    93  	err = c.Connect()
    94  	assert.Nil(t, err)
    95  
    96  	// connecting twice shouldn't do anything
    97  	err = c.Connect()
    98  	assert.Nil(t, err)
    99  
   100  	c.Close()
   101  	c.Output()
   102  }
   103  
   104  func TestConnectAuthFailure(t *testing.T) {
   105  	s := httptest.NewServer(http.HandlerFunc(connect))
   106  	defer s.Close()
   107  
   108  	log := logrus.New()
   109  	log.SetLevel(logrus.DebugLevel)
   110  	u := "ws" + strings.TrimPrefix(s.URL, "http")
   111  	var retries uint64 = 0
   112  	c, err := New(Config{
   113  		APIKey:     "bad",
   114  		Feed:       Feed(u),
   115  		Market:     Market(""),
   116  		Log:        log,
   117  		MaxRetries: &retries,
   118  	})
   119  	assert.NotNil(t, c)
   120  	assert.Nil(t, err)
   121  
   122  	defer func() {
   123  		time.Sleep(100 * time.Millisecond)
   124  		c.Close()
   125  	}()
   126  
   127  	err = c.Connect()
   128  	assert.Nil(t, err)
   129  
   130  	err = <-c.Error()
   131  	assert.NotNil(t, err)
   132  }
   133  
   134  func TestConnectRetryFailure(t *testing.T) {
   135  	s := httptest.NewServer(http.HandlerFunc(connect))
   136  	defer s.Close()
   137  
   138  	log := logrus.New()
   139  	log.SetLevel(logrus.DebugLevel)
   140  	u := "wss" + strings.TrimPrefix(s.URL, "http") // connecting to wss should fail
   141  	var retries uint64 = 1
   142  	c, err := New(Config{
   143  		APIKey:     "bad",
   144  		Feed:       Feed(u),
   145  		Market:     Market(""),
   146  		Log:        log,
   147  		MaxRetries: &retries,
   148  	})
   149  	assert.NotNil(t, c)
   150  	assert.Nil(t, err)
   151  	err = c.Connect()
   152  	assert.NotNil(t, err)
   153  	c.Close()
   154  }