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

     1  package polygonws
     2  
     3  import (
     4  	"encoding/json"
     5  	"testing"
     6  
     7  	"github.com/polygon-io/client-go/websocket/models"
     8  	"github.com/stretchr/testify/assert"
     9  )
    10  
    11  func TestSupportsTopic(t *testing.T) {
    12  	assert.Equal(t, true, Stocks.supports(StocksMinAggs))
    13  	assert.Equal(t, false, Stocks.supports(stocksMax))
    14  	assert.Equal(t, true, Options.supports(OptionsSecAggs))
    15  	assert.Equal(t, false, Options.supports(StocksMinAggs))
    16  	assert.Equal(t, true, Forex.supports(ForexQuotes))
    17  	assert.Equal(t, false, Forex.supports(OptionsQuotes))
    18  	assert.Equal(t, true, Crypto.supports(CryptoL2Book))
    19  	assert.Equal(t, false, Crypto.supports(cryptoMin))
    20  	assert.Equal(t, true, Market("testMarket").supports(StocksImbalances))
    21  }
    22  
    23  func TestGet(t *testing.T) {
    24  	c, err := New(Config{
    25  		APIKey: "test",
    26  		Feed:   RealTime,
    27  		Market: Stocks,
    28  	})
    29  	assert.Nil(t, err)
    30  
    31  	// subscribe to a topic
    32  	minAggs := models.ControlMessage{Action: models.Subscribe, Params: "AM.AAPL"}
    33  	aBytes, _ := json.Marshal(minAggs)
    34  	err = c.Subscribe(StocksMinAggs, "AAPL")
    35  	assert.Nil(t, err)
    36  	msgs := c.subs.get()
    37  	assert.Equal(t, []json.RawMessage{aBytes}, msgs)
    38  
    39  	// unsubscribe from *
    40  	err = c.Unsubscribe(StocksMinAggs, "NFLX", "*")
    41  	assert.Nil(t, err)
    42  	msgs = c.subs.get()
    43  	assert.Equal(t, []json.RawMessage(nil), msgs)
    44  
    45  	// subscribe to another topic
    46  	err = c.Subscribe(StocksTrades, "SNAP")
    47  	assert.Nil(t, err)
    48  	trades := models.ControlMessage{Action: models.Subscribe, Params: "T.SNAP"}
    49  	tBytes, _ := json.Marshal(trades)
    50  	msgs = c.subs.get()
    51  	assert.Equal(t, []json.RawMessage{tBytes}, msgs)
    52  
    53  	// unsubscribe from *
    54  	err = c.Unsubscribe(StocksTrades)
    55  	assert.Nil(t, err)
    56  	msgs = c.subs.get()
    57  	assert.Equal(t, []json.RawMessage(nil), msgs)
    58  }
    59  
    60  func TestGetSub(t *testing.T) {
    61  	submsg, err := getSub(models.Subscribe, StocksMinAggs, "AAPL", "GME", "HOOD")
    62  	assert.Nil(t, err)
    63  	sub, err := json.Marshal(models.ControlMessage{Action: models.Subscribe, Params: "AM.AAPL,AM.GME,AM.HOOD"})
    64  	assert.Nil(t, err)
    65  	assert.Equal(t, json.RawMessage(sub), submsg)
    66  
    67  	unsubmsg, err := getSub(models.Unsubscribe, StocksSecAggs)
    68  	assert.Nil(t, err)
    69  	unsub, err := json.Marshal(models.ControlMessage{Action: models.Unsubscribe, Params: "A.*"})
    70  	assert.Nil(t, err)
    71  	assert.Equal(t, json.RawMessage(unsub), unsubmsg)
    72  }
    73  
    74  func TestSubscriptions(t *testing.T) {
    75  	c, err := New(Config{
    76  		APIKey: "test",
    77  		Feed:   RealTime,
    78  		Market: Stocks,
    79  	})
    80  	assert.Nil(t, err)
    81  
    82  	err = c.Subscribe(StocksMinAggs, "AAPL", "TSLA")
    83  	assert.Nil(t, err)
    84  	_, aapl := c.subs[StocksMinAggs]["AAPL"]
    85  	assert.True(t, aapl)
    86  	_, tsla := c.subs[StocksMinAggs]["TSLA"]
    87  	assert.True(t, tsla)
    88  
    89  	err = c.Unsubscribe(StocksMinAggs, "AAPL", "NFLX")
    90  	assert.Nil(t, err)
    91  	_, aapl = c.subs[StocksMinAggs]["AAPL"]
    92  	assert.False(t, aapl)
    93  
    94  	err = c.Subscribe(StocksMinAggs)
    95  	assert.Nil(t, err)
    96  	_, all := c.subs[StocksMinAggs]["*"]
    97  	assert.True(t, all)
    98  	_, tsla = c.subs[StocksMinAggs]["TSLA"]
    99  	assert.False(t, tsla)
   100  
   101  	err = c.Unsubscribe(StocksMinAggs, "*")
   102  	assert.Nil(t, err)
   103  	_, all = c.subs[StocksMinAggs]["*"]
   104  	assert.False(t, all)
   105  	_, trade := c.subs[StocksTrades]
   106  	assert.False(t, trade)
   107  
   108  	err = c.Unsubscribe(StocksTrades, "RDFN")
   109  	assert.Nil(t, err)
   110  	_, trade = c.subs[StocksTrades]
   111  	assert.False(t, trade)
   112  
   113  	err = c.Subscribe(StocksTrades, "FB")
   114  	assert.Nil(t, err)
   115  	_, fb := c.subs[StocksTrades]["FB"]
   116  	assert.True(t, fb)
   117  
   118  	err = c.Unsubscribe(StocksTrades)
   119  	assert.Nil(t, err)
   120  	_, fb = c.subs[StocksTrades]["FB"]
   121  	assert.False(t, fb)
   122  }