package stream import ( "bytes" "compress/flate" "compress/gzip" "encoding/json" "errors" "fmt" "net" "net/http" "sort" "strconv" "strings" "sync" "testing" "time" "github.com/gorilla/websocket" "github.com/stretchr/testify/assert" "github.com/thrasher-corp/gocryptotrader/config" "github.com/thrasher-corp/gocryptotrader/currency" "github.com/thrasher-corp/gocryptotrader/exchanges/protocol" "github.com/thrasher-corp/gocryptotrader/exchanges/subscription" ) const ( websocketTestURL = "wss://www.bitmex.com/realtime" useProxyTests = false // Disabled by default. Freely available proxy servers that work all the time are difficult to find proxyURL = "http://212.186.171.4:80" // Replace with a usable proxy server ) var dialer websocket.Dialer type testStruct struct { Error error WC WebsocketConnection } type testRequest struct { Event string `json:"event"` RequestID int64 `json:"reqid,omitempty"` Pairs []string `json:"pair"` Subscription testRequestData `json:"subscription,omitempty"` } // testRequestData contains details on WS channel type testRequestData struct { Name string `json:"name,omitempty"` Interval int64 `json:"interval,omitempty"` Depth int64 `json:"depth,omitempty"` } type testResponse struct { RequestID int64 `json:"reqid,omitempty"` } type testSubKey struct { Mood string } var defaultSetup = &WebsocketSetup{ ExchangeConfig: &config.Exchange{ Features: &config.FeaturesConfig{ Enabled: config.FeaturesEnabledConfig{Websocket: true}, }, API: config.APIConfig{ AuthenticatedWebsocketSupport: true, }, WebsocketTrafficTimeout: time.Second * 5, Name: "exchangeName", }, DefaultURL: "testDefaultURL", RunningURL: "wss://testRunningURL", Connector: func() error { return nil }, Subscriber: func(_ []subscription.Subscription) error { return nil }, Unsubscriber: func(_ []subscription.Subscription) error { return nil }, GenerateSubscriptions: func() ([]subscription.Subscription, error) { return []subscription.Subscription{ {Channel: "TestSub"}, {Channel: "TestSub2", Key: "purple"}, {Channel: "TestSub3", Key: testSubKey{"mauve"}}, {Channel: "TestSub4", Key: 42}, }, nil }, Features: &protocol.Features{Subscribe: true, Unsubscribe: true}, } type dodgyConnection struct { WebsocketConnection } // override websocket connection method to produce a wicked terrible error func (d *dodgyConnection) Shutdown() error { return errors.New("cannot shutdown due to some dastardly reason") } // override websocket connection method to produce a wicked terrible error func (d *dodgyConnection) Connect() error { return errors.New("cannot connect due to some dastardly reason") } func TestSetup(t *testing.T) { t.Parallel() var w *Websocket err := w.Setup(nil) if !errors.Is(err, errWebsocketIsNil) { t.Fatalf("received: '%v' but expected: '%v'", err, errWebsocketIsNil) } w = &Websocket{DataHandler: make(chan interface{})} err = w.Setup(nil) if !errors.Is(err, errWebsocketSetupIsNil) { t.Fatalf("received: '%v' but expected: '%v'", err, errWebsocketSetupIsNil) } websocketSetup := &WebsocketSetup{} err = w.Setup(websocketSetup) if !errors.Is(err, errWebsocketAlreadyInitialised) { t.Fatalf("received: '%v' but expected: '%v'", err, errWebsocketAlreadyInitialised) } w.Init = true err = w.Setup(websocketSetup) if !errors.Is(err, errExchangeConfigIsNil) { t.Fatalf("received: '%v' but expected: '%v'", err, errExchangeConfigIsNil) } websocketSetup.ExchangeConfig = &config.Exchange{} err = w.Setup(websocketSetup) if !errors.Is(err, errExchangeConfigNameUnset) { t.Fatalf("received: '%v' but expected: '%v'", err, errExchangeConfigNameUnset) } websocketSetup.ExchangeConfig.Name = "testname" err = w.Setup(websocketSetup) if !errors.Is(err, errWebsocketFeaturesIsUnset) { t.Fatalf("received: '%v' but expected: '%v'", err, errWebsocketFeaturesIsUnset) } websocketSetup.Features = &protocol.Features{} err = w.Setup(websocketSetup) if !errors.Is(err, errConfigFeaturesIsNil) { t.Fatalf("received: '%v' but expected: '%v'", err, errConfigFeaturesIsNil) } websocketSetup.ExchangeConfig.Features = &config.FeaturesConfig{} err = w.Setup(websocketSetup) if !errors.Is(err, errWebsocketConnectorUnset) { t.Fatalf("received: '%v' but expected: '%v'", err, errWebsocketConnectorUnset) } websocketSetup.Connector = func() error { return nil } err = w.Setup(websocketSetup) if !errors.Is(err, errWebsocketSubscriberUnset) { t.Fatalf("received: '%v' but expected: '%v'", err, errWebsocketSubscriberUnset) } websocketSetup.Subscriber = func([]subscription.Subscription) error { return nil } websocketSetup.Features.Unsubscribe = true err = w.Setup(websocketSetup) if !errors.Is(err, errWebsocketUnsubscriberUnset) { t.Fatalf("received: '%v' but expected: '%v'", err, errWebsocketUnsubscriberUnset) } websocketSetup.Unsubscriber = func([]subscription.Subscription) error { return nil } err = w.Setup(websocketSetup) if !errors.Is(err, errWebsocketSubscriptionsGeneratorUnset) { t.Fatalf("received: '%v' but expected: '%v'", err, errWebsocketSubscriptionsGeneratorUnset) } websocketSetup.GenerateSubscriptions = func() ([]subscription.Subscription, error) { return nil, nil } err = w.Setup(websocketSetup) if !errors.Is(err, errDefaultURLIsEmpty) { t.Fatalf("received: '%v' but expected: '%v'", err, errDefaultURLIsEmpty) } websocketSetup.DefaultURL = "test" err = w.Setup(websocketSetup) if !errors.Is(err, errRunningURLIsEmpty) { t.Fatalf("received: '%v' but expected: '%v'", err, errRunningURLIsEmpty) } websocketSetup.RunningURL = "http://www.google.com" err = w.Setup(websocketSetup) if !errors.Is(err, errInvalidWebsocketURL) { t.Fatalf("received: '%v' but expected: '%v'", err, errInvalidWebsocketURL) } websocketSetup.RunningURL = "wss://www.google.com" websocketSetup.RunningURLAuth = "http://www.google.com" err = w.Setup(websocketSetup) if !errors.Is(err, errInvalidWebsocketURL) { t.Fatalf("received: '%v' but expected: '%v'", err, errInvalidWebsocketURL) } websocketSetup.RunningURLAuth = "wss://www.google.com" err = w.Setup(websocketSetup) if !errors.Is(err, errInvalidTrafficTimeout) { t.Fatalf("received: '%v' but expected: '%v'", err, errInvalidTrafficTimeout) } websocketSetup.ExchangeConfig.WebsocketTrafficTimeout = time.Minute err = w.Setup(websocketSetup) if !errors.Is(err, nil) { t.Fatalf("received: %v but expected: %v", err, nil) } } func TestTrafficMonitorTimeout(t *testing.T) { t.Parallel() ws := *New() if err := ws.Setup(defaultSetup); err != nil { t.Fatal(err) } ws.trafficTimeout = time.Second * 2 ws.ShutdownC = make(chan struct{}) ws.trafficMonitor() if !ws.IsTrafficMonitorRunning() { t.Fatal("traffic monitor should be running") } // Deploy traffic alert ws.TrafficAlert <- struct{}{} // try to add another traffic monitor ws.trafficMonitor() if !ws.IsTrafficMonitorRunning() { t.Fatal("traffic monitor should be running") } // prevent shutdown routine ws.setConnectedStatus(false) // await timeout closure ws.Wg.Wait() if ws.IsTrafficMonitorRunning() { t.Error("should be dead") } } func TestIsDisconnectionError(t *testing.T) { t.Parallel() isADisconnectionError := IsDisconnectionError(errors.New("errorText")) if isADisconnectionError { t.Error("Its not") } isADisconnectionError = IsDisconnectionError(&websocket.CloseError{ Code: 1006, Text: "errorText", }) if !isADisconnectionError { t.Error("It is") } isADisconnectionError = IsDisconnectionError(&net.OpError{ Err: errClosedConnection, }) if isADisconnectionError { t.Error("It's not") } isADisconnectionError = IsDisconnectionError(&net.OpError{ Err: errors.New("errText"), }) if !isADisconnectionError { t.Error("It is") } } func TestConnectionMessageErrors(t *testing.T) { t.Parallel() var wsWrong = &Websocket{} err := wsWrong.Connect() if err == nil { t.Fatal("error cannot be nil") } wsWrong.connector = func() error { return nil } err = wsWrong.Connect() if err == nil { t.Fatal("error cannot be nil") } wsWrong.setEnabled(true) wsWrong.setConnectingStatus(true) wsWrong.Wg = &sync.WaitGroup{} err = wsWrong.Connect() if err == nil { t.Fatal("error cannot be nil") } wsWrong.setConnectedStatus(false) wsWrong.connector = func() error { return errors.New("edge case error of dooooooom") } err = wsWrong.Connect() if err == nil { t.Fatal("error cannot be nil") } ws := *New() err = ws.Setup(defaultSetup) if err != nil { t.Fatal(err) } ws.trafficTimeout = time.Minute ws.connector = func() error { return nil } err = ws.Connect() if err != nil { t.Fatal(err) } ws.TrafficAlert <- struct{}{} timer := time.NewTimer(900 * time.Millisecond) ws.ReadMessageErrors <- errors.New("errorText") select { case err := <-ws.ToRoutine: errText, ok := err.(error) if !ok { t.Error("unable to type assert error") } else if errText.Error() != "errorText" { t.Errorf("Expected 'errorText', received %v", err) } case <-timer.C: t.Error("Timeout waiting for datahandler to receive error") } ws.ReadMessageErrors <- &websocket.CloseError{ Code: 1006, Text: "errorText", } outer: for { select { case err := <-ws.ToRoutine: if _, ok := err.(*websocket.CloseError); !ok { t.Errorf("Error is not a disconnection error: %v", err) } case <-timer.C: break outer } } } func TestWebsocket(t *testing.T) { t.Parallel() wsInit := Websocket{} err := wsInit.Setup(&WebsocketSetup{ ExchangeConfig: &config.Exchange{ Features: &config.FeaturesConfig{ Enabled: config.FeaturesEnabledConfig{Websocket: true}, }, Name: "test", }, }) if !errors.Is(err, errWebsocketAlreadyInitialised) { t.Fatalf("received: '%v' but expected: '%v'", err, errWebsocketAlreadyInitialised) } ws := *New() err = ws.SetProxyAddress("garbagio") if err == nil { t.Error("error cannot be nil") } ws.Conn = &WebsocketConnection{} ws.AuthConn = &WebsocketConnection{} ws.setEnabled(true) err = ws.SetProxyAddress("https://192.168.0.1:1337") if err == nil { t.Error("error cannot be nil") } ws.setConnectedStatus(true) ws.ShutdownC = make(chan struct{}) ws.Wg = &sync.WaitGroup{} err = ws.SetProxyAddress("https://192.168.0.1:1336") if err == nil { t.Error("SetProxyAddress", err) } err = ws.SetProxyAddress("https://192.168.0.1:1336") if err == nil { t.Error("SetProxyAddress", err) } ws.setEnabled(false) // removing proxy err = ws.SetProxyAddress("") if err != nil { t.Error(err) } // reinstate proxy err = ws.SetProxyAddress("http://localhost:1337") if err != nil { t.Error(err) } // conflict proxy err = ws.SetProxyAddress("http://localhost:1337") if err == nil { t.Error("error cannot be nil") } err = ws.Setup(defaultSetup) if err != nil { t.Fatal(err) } if ws.GetName() != "exchangeName" { t.Error("WebsocketSetup") } if !ws.IsEnabled() { t.Error("WebsocketSetup") } ws.setEnabled(false) if ws.IsEnabled() { t.Error("WebsocketSetup") } ws.setEnabled(true) if !ws.IsEnabled() { t.Error("WebsocketSetup") } if ws.GetProxyAddress() != "http://localhost:1337" { t.Error("WebsocketSetup") } if ws.GetWebsocketURL() != "wss://testRunningURL" { t.Error("WebsocketSetup") } if ws.trafficTimeout != time.Second*5 { t.Error("WebsocketSetup") } // -- Not connected shutdown err = ws.Shutdown() if err == nil { t.Fatal("should not be connected to able to shut down") } ws.setConnectedStatus(true) ws.Conn = &dodgyConnection{} err = ws.Shutdown() if err == nil { t.Fatal("error cannot be nil") } ws.Conn = &WebsocketConnection{} ws.setConnectedStatus(true) ws.AuthConn = &dodgyConnection{} err = ws.Shutdown() if err == nil { t.Fatal("error cannot be nil ") } ws.AuthConn = &WebsocketConnection{} ws.setConnectedStatus(false) // -- Normal connect err = ws.Connect() if err != nil { t.Fatal("WebsocketSetup", err) } ws.defaultURL = "ws://demos.kaazing.com/echo" ws.defaultURLAuth = "ws://demos.kaazing.com/echo" err = ws.SetWebsocketURL("", false, false) if err != nil { t.Fatal(err) } err = ws.SetWebsocketURL("ws://demos.kaazing.com/echo", false, false) if err != nil { t.Fatal(err) } err = ws.SetWebsocketURL("", true, false) if err != nil { t.Fatal(err) } err = ws.SetWebsocketURL("ws://demos.kaazing.com/echo", true, false) if err != nil { t.Fatal(err) } // Attempt reconnect err = ws.SetWebsocketURL("ws://demos.kaazing.com/echo", true, true) if err != nil { t.Fatal(err) } // -- initiate the reconnect which is usually handled by connection monitor err = ws.Connect() if err != nil { t.Fatal(err) } err = ws.Connect() if err == nil { t.Fatal("should already be connected") } // -- Normal shutdown err = ws.Shutdown() if err != nil { t.Fatal("WebsocketSetup", err) } ws.Wg.Wait() } // TestSubscribe logic test func TestSubscribeUnsubscribe(t *testing.T) { t.Parallel() ws := *New() assert.NoError(t, ws.Setup(defaultSetup), "WS Setup should not error") fnSub := func(subs []subscription.Subscription) error { ws.AddSuccessfulSubscriptions(subs...) return nil } fnUnsub := func(unsubs []subscription.Subscription) error { ws.RemoveSubscriptions(unsubs...) return nil } ws.Subscriber = fnSub ws.Unsubscriber = fnUnsub subs, err := ws.GenerateSubs() assert.NoError(t, err, "Generating test subscriptions should not error") assert.ErrorIs(t, ws.UnsubscribeChannels(nil), errNoSubscriptionsSupplied, "Unsubscribing from nil should error") assert.ErrorIs(t, ws.UnsubscribeChannels(subs), ErrSubscriptionNotFound, "Unsubscribing should error when not subscribed") assert.Nil(t, ws.GetSubscription(42), "GetSubscription on empty internal map should return") assert.NoError(t, ws.SubscribeToChannels(subs), "Basic Subscribing should not error") assert.Len(t, ws.GetSubscriptions(), 4, "Should have 4 subscriptions") byDefKey := ws.GetSubscription(subscription.DefaultKey{Channel: "TestSub"}) if assert.NotNil(t, byDefKey, "GetSubscription by default key should find a channel") { assert.Equal(t, "TestSub", byDefKey.Channel, "GetSubscription by default key should return a pointer a copy of the right channel") assert.NotSame(t, byDefKey, ws.subscriptions["TestSub"], "GetSubscription returns a fresh pointer") } if assert.NotNil(t, ws.GetSubscription("purple"), "GetSubscription by string key should find a channel") { assert.Equal(t, "TestSub2", ws.GetSubscription("purple").Channel, "GetSubscription by string key should return a pointer a copy of the right channel") } if assert.NotNil(t, ws.GetSubscription(testSubKey{"mauve"}), "GetSubscription by type key should find a channel") { assert.Equal(t, "TestSub3", ws.GetSubscription(testSubKey{"mauve"}).Channel, "GetSubscription by type key should return a pointer a copy of the right channel") } if assert.NotNil(t, ws.GetSubscription(42), "GetSubscription by int key should find a channel") { assert.Equal(t, "TestSub4", ws.GetSubscription(42).Channel, "GetSubscription by int key should return a pointer a copy of the right channel") } assert.Nil(t, ws.GetSubscription(nil), "GetSubscription by nil should return nil") assert.Nil(t, ws.GetSubscription(45), "GetSubscription by invalid key should return nil") assert.ErrorIs(t, ws.SubscribeToChannels(subs), errChannelAlreadySubscribed, "Subscribe should error when already subscribed") assert.ErrorIs(t, ws.SubscribeToChannels(nil), errNoSubscriptionsSupplied, "Subscribe to nil should error") assert.NoError(t, ws.UnsubscribeChannels(subs), "Unsubscribing should not error") } // TestResubscribe tests Resubscribing to existing subscriptions func TestResubscribe(t *testing.T) { t.Parallel() ws := *New() wackedOutSetup := *defaultSetup wackedOutSetup.MaxWebsocketSubscriptionsPerConnection = -1 err := ws.Setup(&wackedOutSetup) assert.ErrorIs(t, err, errInvalidMaxSubscriptions, "Invalid MaxWebsocketSubscriptionsPerConnection should error") err = ws.Setup(defaultSetup) assert.NoError(t, err, "WS Setup should not error") fnSub := func(subs []subscription.Subscription) error { ws.AddSuccessfulSubscriptions(subs...) return nil } fnUnsub := func(unsubs []subscription.Subscription) error { ws.RemoveSubscriptions(unsubs...) return nil } ws.Subscriber = fnSub ws.Unsubscriber = fnUnsub channel := []subscription.Subscription{{Channel: "resubTest"}} assert.ErrorIs(t, ws.ResubscribeToChannel(&channel[0]), ErrSubscriptionNotFound, "Resubscribe should error when channel isn't subscribed yet") assert.NoError(t, ws.SubscribeToChannels(channel), "Subscribe should not error") assert.NoError(t, ws.ResubscribeToChannel(&channel[0]), "Resubscribe should not error now the channel is subscribed") } // TestSubscriptionState tests Subscription state changes func TestSubscriptionState(t *testing.T) { t.Parallel() ws := New() c := &subscription.Subscription{Key: 42, Channel: "Gophers", State: subscription.SubscribingState} assert.ErrorIs(t, ws.SetSubscriptionState(c, subscription.UnsubscribingState), ErrSubscriptionNotFound, "Setting an imaginary sub should error") assert.NoError(t, ws.AddSubscription(c), "Adding first subscription should not error") found := ws.GetSubscription(42) assert.NotNil(t, found, "Should find the subscription") assert.Equal(t, subscription.SubscribingState, found.State, "Subscription should be Subscribing") assert.ErrorIs(t, ws.AddSubscription(c), ErrSubscribedAlready, "Adding an already existing sub should error") assert.ErrorIs(t, ws.SetSubscriptionState(c, subscription.SubscribingState), ErrChannelInStateAlready, "Setting Same state should error") assert.ErrorIs(t, ws.SetSubscriptionState(c, subscription.UnsubscribingState+1), errInvalidChannelState, "Setting an invalid state should error") ws.AddSuccessfulSubscriptions(*c) found = ws.GetSubscription(42) assert.NotNil(t, found, "Should find the subscription") assert.Equal(t, found.State, subscription.SubscribedState, "Subscription should be subscribed state") assert.NoError(t, ws.SetSubscriptionState(c, subscription.UnsubscribingState), "Setting Unsub state should not error") found = ws.GetSubscription(42) assert.Equal(t, found.State, subscription.UnsubscribingState, "Subscription should be unsubscribing state") } // TestRemoveSubscriptions tests removing a subscription func TestRemoveSubscriptions(t *testing.T) { t.Parallel() ws := New() c := &subscription.Subscription{Key: 42, Channel: "Unite!"} assert.NoError(t, ws.AddSubscription(c), "Adding first subscription should not error") assert.NotNil(t, ws.GetSubscription(42), "Added subscription should be findable") ws.RemoveSubscriptions(*c) assert.Nil(t, ws.GetSubscription(42), "Remove should have removed the sub") } // TestConnectionMonitorNoConnection logic test func TestConnectionMonitorNoConnection(t *testing.T) { t.Parallel() ws := *New() ws.connectionMonitorDelay = 500 ws.DataHandler = make(chan interface{}, 1) ws.ShutdownC = make(chan struct{}, 1) ws.exchangeName = "hello" ws.Wg = &sync.WaitGroup{} ws.enabled = true err := ws.connectionMonitor() if !errors.Is(err, nil) { t.Fatalf("received: %v, but expected: %v", err, nil) } if !ws.IsConnectionMonitorRunning() { t.Fatal("Should not have exited") } err = ws.connectionMonitor() if !errors.Is(err, errAlreadyRunning) { t.Fatalf("received: %v, but expected: %v", err, errAlreadyRunning) } } // TestGetSubscription logic test func TestGetSubscription(t *testing.T) { t.Parallel() assert.Nil(t, (*Websocket).GetSubscription(nil, "imaginary"), "GetSubscription on a nil Websocket should return nil") assert.Nil(t, (&Websocket{}).GetSubscription("empty"), "GetSubscription on a Websocket with no sub map should return nil") w := Websocket{ subscriptions: subscriptionMap{ 42: { Channel: "hello3", }, }, } assert.Nil(t, w.GetSubscription(43), "GetSubscription with an invalid key should return nil") c := w.GetSubscription(42) if assert.NotNil(t, c, "GetSubscription with an valid key should return a channel") { assert.Equal(t, "hello3", c.Channel, "GetSubscription should return the correct channel details") } } // TestGetSubscriptions logic test func TestGetSubscriptions(t *testing.T) { t.Parallel() w := Websocket{ subscriptions: subscriptionMap{ 42: { Channel: "hello3", }, }, } assert.Equal(t, "hello3", w.GetSubscriptions()[0].Channel, "GetSubscriptions should return the correct channel details") } // TestSetCanUseAuthenticatedEndpoints logic test func TestSetCanUseAuthenticatedEndpoints(t *testing.T) { t.Parallel() ws := *New() result := ws.CanUseAuthenticatedEndpoints() if result { t.Error("expected `canUseAuthenticatedEndpoints` to be false") } ws.SetCanUseAuthenticatedEndpoints(true) result = ws.CanUseAuthenticatedEndpoints() if !result { t.Error("expected `canUseAuthenticatedEndpoints` to be true") } } // TestDial logic test func TestDial(t *testing.T) { t.Parallel() var testCases = []testStruct{ {Error: nil, WC: WebsocketConnection{ ExchangeName: "test1", Verbose: true, URL: websocketTestURL, RateLimit: 10, ResponseMaxLimit: 7000000000, }, }, {Error: errors.New(" Error: malformed ws or wss URL"), WC: WebsocketConnection{ ExchangeName: "test2", Verbose: true, URL: "", ResponseMaxLimit: 7000000000, }, }, {Error: nil, WC: WebsocketConnection{ ExchangeName: "test3", Verbose: true, URL: websocketTestURL, ProxyURL: proxyURL, ResponseMaxLimit: 7000000000, }, }, } for i := range testCases { testData := &testCases[i] t.Run(testData.WC.ExchangeName, func(t *testing.T) { t.Parallel() if testData.WC.ProxyURL != "" && !useProxyTests { t.Skip("Proxy testing not enabled, skipping") } err := testData.WC.Dial(&dialer, http.Header{}) if err != nil { if testData.Error != nil && strings.Contains(err.Error(), testData.Error.Error()) { return } t.Fatal(err) } }) } } // TestSendMessage logic test func TestSendMessage(t *testing.T) { t.Parallel() var testCases = []testStruct{ {Error: nil, WC: WebsocketConnection{ ExchangeName: "test1", Verbose: true, URL: websocketTestURL, RateLimit: 10, ResponseMaxLimit: 7000000000, }, }, {Error: errors.New(" Error: malformed ws or wss URL"), WC: WebsocketConnection{ ExchangeName: "test2", Verbose: true, URL: "", ResponseMaxLimit: 7000000000, }, }, {Error: nil, WC: WebsocketConnection{ ExchangeName: "test3", Verbose: true, URL: websocketTestURL, ProxyURL: proxyURL, ResponseMaxLimit: 7000000000, }, }, } for i := range testCases { testData := &testCases[i] t.Run(testData.WC.ExchangeName, func(t *testing.T) { t.Parallel() if testData.WC.ProxyURL != "" && !useProxyTests { t.Skip("Proxy testing not enabled, skipping") } err := testData.WC.Dial(&dialer, http.Header{}) if err != nil { if testData.Error != nil && strings.Contains(err.Error(), testData.Error.Error()) { return } t.Fatal(err) } err = testData.WC.SendJSONMessage(Ping) if err != nil { t.Error(err) } err = testData.WC.SendRawMessage(websocket.TextMessage, []byte(Ping)) if err != nil { t.Error(err) } }) } } // TestSendMessageWithResponse logic test func TestSendMessageWithResponse(t *testing.T) { t.Parallel() wc := &WebsocketConnection{ Verbose: true, URL: "wss://ws.kraken.com", ResponseMaxLimit: time.Second * 5, Match: NewMatch(), } if wc.ProxyURL != "" && !useProxyTests { t.Skip("Proxy testing not enabled, skipping") } err := wc.Dial(&dialer, http.Header{}) if err != nil { t.Fatal(err) } go readMessages(t, wc) request := testRequest{ Event: "subscribe", Pairs: []string{currency.NewPairWithDelimiter("XBT", "USD", "/").String()}, Subscription: testRequestData{ Name: "ticker", }, RequestID: wc.GenerateMessageID(false), } _, err = wc.SendMessageReturnResponse(request.RequestID, request) if err != nil { t.Error(err) } } type reporter struct { name string msg []byte t time.Duration } func (r *reporter) Latency(name string, message []byte, t time.Duration) { r.name = name r.msg = message r.t = t } // readMessages helper func func readMessages(t *testing.T, wc *WebsocketConnection) { t.Helper() timer := time.NewTimer(20 * time.Second) for { select { case <-timer.C: return default: resp := wc.ReadMessage() if resp.Raw == nil { t.Error("connection has closed") return } var incoming testResponse err := json.Unmarshal(resp.Raw, &incoming) if err != nil { t.Error(err) return } if incoming.RequestID > 0 { wc.Match.IncomingWithData(incoming.RequestID, resp.Raw) return } } } } // TestSetupPingHandler logic test func TestSetupPingHandler(t *testing.T) { t.Parallel() wc := &WebsocketConnection{ URL: websocketTestURL, ResponseMaxLimit: time.Second * 5, Match: NewMatch(), Wg: &sync.WaitGroup{}, } if wc.ProxyURL != "" && !useProxyTests { t.Skip("Proxy testing not enabled, skipping") } wc.ShutdownC = make(chan struct{}) err := wc.Dial(&dialer, http.Header{}) if err != nil { t.Fatal(err) } wc.SetupPingHandler(PingHandler{ UseGorillaHandler: true, MessageType: websocket.PingMessage, Delay: 100, }) err = wc.Connection.Close() if err != nil { t.Error(err) } err = wc.Dial(&dialer, http.Header{}) if err != nil { t.Fatal(err) } wc.SetupPingHandler(PingHandler{ MessageType: websocket.TextMessage, Message: []byte(Ping), Delay: 200, }) time.Sleep(time.Millisecond * 201) close(wc.ShutdownC) wc.Wg.Wait() } // TestParseBinaryResponse logic test func TestParseBinaryResponse(t *testing.T) { t.Parallel() wc := &WebsocketConnection{ URL: websocketTestURL, ResponseMaxLimit: time.Second * 5, Match: NewMatch(), } var b bytes.Buffer w := gzip.NewWriter(&b) _, err := w.Write([]byte("hello")) if err != nil { t.Error(err) } err = w.Close() if err != nil { t.Error(err) } var resp []byte resp, err = wc.parseBinaryResponse(b.Bytes()) if err != nil { t.Error(err) } if !strings.EqualFold(string(resp), "hello") { t.Errorf("GZip conversion failed. Received: '%v', Expected: 'hello'", string(resp)) } var b2 bytes.Buffer w2, err2 := flate.NewWriter(&b2, 1) if err2 != nil { t.Error(err2) } _, err2 = w2.Write([]byte("hello")) if err2 != nil { t.Error(err) } err2 = w2.Close() if err2 != nil { t.Error(err) } resp2, err3 := wc.parseBinaryResponse(b2.Bytes()) if err3 != nil { t.Error(err3) } if !strings.EqualFold(string(resp2), "hello") { t.Errorf("Deflate conversion failed. Received: '%v', Expected: 'hello'", string(resp2)) } _, err4 := wc.parseBinaryResponse([]byte{}) if err4 == nil || err4.Error() != "unexpected EOF" { t.Error("Expected error 'unexpected EOF'") } } // TestCanUseAuthenticatedWebsocketForWrapper logic test func TestCanUseAuthenticatedWebsocketForWrapper(t *testing.T) { t.Parallel() ws := &Websocket{} resp := ws.CanUseAuthenticatedWebsocketForWrapper() if resp { t.Error("Expected false, `connected` is false") } ws.setConnectedStatus(true) resp = ws.CanUseAuthenticatedWebsocketForWrapper() if resp { t.Error("Expected false, `connected` is true and `CanUseAuthenticatedEndpoints` is false") } ws.canUseAuthenticatedEndpoints = true resp = ws.CanUseAuthenticatedWebsocketForWrapper() if !resp { t.Error("Expected true, `connected` and `CanUseAuthenticatedEndpoints` is true") } } func TestGenerateMessageID(t *testing.T) { t.Parallel() wc := WebsocketConnection{} var id int64 for i := 0; i < 10; i++ { newID := wc.GenerateMessageID(true) if id == newID { t.Fatal("ID generation is not unique") } id = newID } } // BenchmarkGenerateMessageID-8 2850018 408 ns/op 56 B/op 4 allocs/op func BenchmarkGenerateMessageID_High(b *testing.B) { wc := WebsocketConnection{} for i := 0; i < b.N; i++ { _ = wc.GenerateMessageID(true) } } // BenchmarkGenerateMessageID_Low-8 2591596 447 ns/op 56 B/op 4 allocs/op func BenchmarkGenerateMessageID_Low(b *testing.B) { wc := WebsocketConnection{} for i := 0; i < b.N; i++ { _ = wc.GenerateMessageID(false) } } func TestCheckWebsocketURL(t *testing.T) { err := checkWebsocketURL("") if err == nil { t.Fatal("error cannot be nil") } err = checkWebsocketURL("wowowow:wowowowo") if err == nil { t.Fatal("error cannot be nil") } err = checkWebsocketURL("://") if err == nil { t.Fatal("error cannot be nil") } err = checkWebsocketURL("http://www.google.com") if err == nil { t.Fatal("error cannot be nil") } err = checkWebsocketURL("wss://websocketconnection.place") if err != nil { t.Fatal(err) } err = checkWebsocketURL("ws://websocketconnection.place") if err != nil { t.Fatal(err) } } func TestGetChannelDifference(t *testing.T) { t.Parallel() web := Websocket{} newChans := []subscription.Subscription{ { Channel: "Test1", }, { Channel: "Test2", }, { Channel: "Test3", }, } subs, unsubs := web.GetChannelDifference(newChans) assert.Len(t, subs, 3, "Should get the correct number of subs") assert.Len(t, unsubs, 0, "Should get the correct number of unsubs") web.AddSuccessfulSubscriptions(subs...) flushedSubs := []subscription.Subscription{ { Channel: "Test2", }, } subs, unsubs = web.GetChannelDifference(flushedSubs) assert.Len(t, subs, 0, "Should get the correct number of subs") assert.Len(t, unsubs, 2, "Should get the correct number of unsubs") flushedSubs = []subscription.Subscription{ { Channel: "Test2", }, { Channel: "Test4", }, } subs, unsubs = web.GetChannelDifference(flushedSubs) if assert.Len(t, subs, 1, "Should get the correct number of subs") { assert.Equal(t, subs[0].Channel, "Test4", "Should subscribe to the right channel") } if assert.Len(t, unsubs, 2, "Should get the correct number of unsubs") { sort.Slice(unsubs, func(i, j int) bool { return unsubs[i].Channel <= unsubs[j].Channel }) assert.Equal(t, unsubs[0].Channel, "Test1", "Should unsubscribe from the right channels") assert.Equal(t, unsubs[1].Channel, "Test3", "Should unsubscribe from the right channels") } } // GenSubs defines a theoretical exchange with pair management type GenSubs struct { EnabledPairs currency.Pairs subscribos []subscription.Subscription unsubscribos []subscription.Subscription } // generateSubs default subs created from the enabled pairs list func (g *GenSubs) generateSubs() ([]subscription.Subscription, error) { superduperchannelsubs := make([]subscription.Subscription, len(g.EnabledPairs)) for i := range g.EnabledPairs { superduperchannelsubs[i] = subscription.Subscription{ Channel: "TEST:" + strconv.FormatInt(int64(i), 10), Pair: g.EnabledPairs[i], } } return superduperchannelsubs, nil } func (g *GenSubs) SUBME(subs []subscription.Subscription) error { if len(subs) == 0 { return errors.New("WOW") } g.subscribos = subs return nil } func (g *GenSubs) UNSUBME(unsubs []subscription.Subscription) error { if len(unsubs) == 0 { return errors.New("WOW") } g.unsubscribos = unsubs return nil } // sneaky connect func func connect() error { return nil } func TestFlushChannels(t *testing.T) { t.Parallel() // Enabled pairs/setup system newgen := GenSubs{EnabledPairs: []currency.Pair{ currency.NewPair(currency.BTC, currency.AUD), currency.NewPair(currency.BTC, currency.USDT), }} dodgyWs := Websocket{} err := dodgyWs.FlushChannels() if err == nil { t.Fatal("error cannot be nil") } dodgyWs.setEnabled(true) err = dodgyWs.FlushChannels() if err == nil { t.Fatal("error cannot be nil") } web := Websocket{ enabled: true, connected: true, connector: connect, ShutdownC: make(chan struct{}), Subscriber: newgen.SUBME, Unsubscriber: newgen.UNSUBME, Wg: new(sync.WaitGroup), features: &protocol.Features{ // No features }, trafficTimeout: time.Second * 30, // Added for when we utilise connect() // in FlushChannels() so the traffic monitor doesn't time out and turn // this to an unconnected state } problemFunc := func() ([]subscription.Subscription, error) { return nil, errors.New("problems") } noSub := func() ([]subscription.Subscription, error) { return nil, nil } // Disable pair and flush system newgen.EnabledPairs = []currency.Pair{ currency.NewPair(currency.BTC, currency.AUD)} web.GenerateSubs = func() ([]subscription.Subscription, error) { return []subscription.Subscription{{Channel: "test"}}, nil } err = web.FlushChannels() if err != nil { t.Fatal(err) } web.features.FullPayloadSubscribe = true web.GenerateSubs = problemFunc err = web.FlushChannels() // error on full subscribeToChannels if err == nil { t.Fatal("error cannot be nil") } web.GenerateSubs = noSub err = web.FlushChannels() // No subs to sub if err != nil { t.Fatal(err) } web.GenerateSubs = newgen.generateSubs subs, err := web.GenerateSubs() if err != nil { t.Fatal(err) } web.AddSuccessfulSubscriptions(subs...) err = web.FlushChannels() if err != nil { t.Fatal(err) } web.features.FullPayloadSubscribe = false web.features.Subscribe = true web.GenerateSubs = problemFunc err = web.FlushChannels() if err == nil { t.Fatal("error cannot be nil") } web.GenerateSubs = newgen.generateSubs err = web.FlushChannels() if err != nil { t.Fatal(err) } web.subscriptionMutex.Lock() web.subscriptions = subscriptionMap{ 41: { Key: 41, Channel: "match channel", Pair: currency.NewPair(currency.BTC, currency.AUD), }, 42: { Key: 42, Channel: "unsub channel", Pair: currency.NewPair(currency.THETA, currency.USDT), }, } web.subscriptionMutex.Unlock() err = web.FlushChannels() if err != nil { t.Fatal(err) } err = web.FlushChannels() if err != nil { t.Fatal(err) } web.setConnectedStatus(true) web.features.Unsubscribe = true err = web.FlushChannels() if err != nil { t.Fatal(err) } } func TestDisable(t *testing.T) { t.Parallel() web := Websocket{ enabled: true, connected: true, ShutdownC: make(chan struct{}), } err := web.Disable() if err != nil { t.Fatal(err) } err = web.Disable() if err == nil { t.Fatal("should already be disabled") } } func TestEnable(t *testing.T) { t.Parallel() web := Websocket{ connector: connect, Wg: new(sync.WaitGroup), ShutdownC: make(chan struct{}), GenerateSubs: func() ([]subscription.Subscription, error) { return []subscription.Subscription{{Channel: "test"}}, nil }, Subscriber: func(cs []subscription.Subscription) error { return nil }, } err := web.Enable() if err != nil { t.Fatal(err) } err = web.Enable() if err == nil { t.Fatal("should already be enabled") } fmt.Print() } func TestSetupNewConnection(t *testing.T) { t.Parallel() var nonsenseWebsock *Websocket err := nonsenseWebsock.SetupNewConnection(ConnectionSetup{URL: "urlstring"}) if err == nil { t.Fatal("error cannot be nil") } nonsenseWebsock = &Websocket{} err = nonsenseWebsock.SetupNewConnection(ConnectionSetup{URL: "urlstring"}) if err == nil { t.Fatal("error cannot be nil") } nonsenseWebsock = &Websocket{exchangeName: "test"} err = nonsenseWebsock.SetupNewConnection(ConnectionSetup{URL: "urlstring"}) if err == nil { t.Fatal("error cannot be nil") } nonsenseWebsock.TrafficAlert = make(chan struct{}) err = nonsenseWebsock.SetupNewConnection(ConnectionSetup{URL: "urlstring"}) if err == nil { t.Fatal("error cannot be nil") } web := Websocket{ connector: connect, Wg: new(sync.WaitGroup), ShutdownC: make(chan struct{}), Init: true, TrafficAlert: make(chan struct{}), ReadMessageErrors: make(chan error), DataHandler: make(chan interface{}), } err = web.Setup(defaultSetup) if err != nil { t.Fatal(err) } err = web.SetupNewConnection(ConnectionSetup{}) if err == nil { t.Fatal("error cannot be nil") } err = web.SetupNewConnection(ConnectionSetup{URL: "urlstring"}) if err != nil { t.Fatal(err) } err = web.SetupNewConnection(ConnectionSetup{URL: "urlstring", Authenticated: true}) if err != nil { t.Fatal(err) } } func TestWebsocketConnectionShutdown(t *testing.T) { t.Parallel() wc := WebsocketConnection{} err := wc.Shutdown() if err != nil { t.Fatal(err) } err = wc.Dial(&websocket.Dialer{}, nil) if err == nil { t.Fatal("error cannot be nil") } wc.URL = websocketTestURL err = wc.Dial(&websocket.Dialer{}, nil) if err != nil { t.Fatal(err) } err = wc.Shutdown() if err != nil { t.Fatal(err) } } // TestLatency logic test func TestLatency(t *testing.T) { t.Parallel() r := &reporter{} exch := "Kraken" wc := &WebsocketConnection{ ExchangeName: exch, Verbose: true, URL: "wss://ws.kraken.com", ResponseMaxLimit: time.Second * 5, Match: NewMatch(), Reporter: r, } if wc.ProxyURL != "" && !useProxyTests { t.Skip("Proxy testing not enabled, skipping") } err := wc.Dial(&dialer, http.Header{}) if err != nil { t.Fatal(err) } go readMessages(t, wc) request := testRequest{ Event: "subscribe", Pairs: []string{currency.NewPairWithDelimiter("XBT", "USD", "/").String()}, Subscription: testRequestData{ Name: "ticker", }, RequestID: wc.GenerateMessageID(false), } _, err = wc.SendMessageReturnResponse(request.RequestID, request) if err != nil { t.Error(err) } if r.t == 0 { t.Error("expected a nonzero duration, got zero") } if r.name != exch { t.Errorf("expected %v, got %v", exch, r.name) } } func TestCheckSubscriptions(t *testing.T) { t.Parallel() ws := Websocket{} err := ws.checkSubscriptions(nil) if !errors.Is(err, errNoSubscriptionsSupplied) { t.Fatalf("received: %v, but expected: %v", err, errNoSubscriptionsSupplied) } ws.MaxSubscriptionsPerConnection = 1 err = ws.checkSubscriptions([]subscription.Subscription{{}, {}}) if !errors.Is(err, errSubscriptionsExceedsLimit) { t.Fatalf("received: %v, but expected: %v", err, errSubscriptionsExceedsLimit) } ws.MaxSubscriptionsPerConnection = 2 ws.subscriptions = subscriptionMap{42: {Key: 42, Channel: "test"}} err = ws.checkSubscriptions([]subscription.Subscription{{Key: 42, Channel: "test"}}) if !errors.Is(err, errChannelAlreadySubscribed) { t.Fatalf("received: %v, but expected: %v", err, errChannelAlreadySubscribed) } err = ws.checkSubscriptions([]subscription.Subscription{{}}) if !errors.Is(err, nil) { t.Fatalf("received: %v, but expected: %v", err, nil) } }