WS test improvements (#303)

* Improves okgroup and kraken websocket tests so there is less overhead required * Fixes declaration shadowing * Fixes duplication * Switches error to fatal for wsconnection problems in test.
2026-05-23 23:16:49 +00:00 · 2019-05-20 10:32:39 +10:00
parent 05a7ae70ba
commit 08542b6ef7
3 changed files with 80 additions and 136 deletions
--- a/exchanges/kraken/kraken_test.go
+++ b/exchanges/kraken/kraken_test.go
@@ -651,9 +651,6 @@ func TestOrderbookBufferReset(t *testing.T) {
 	if !k.Websocket.IsEnabled() {
 		t.Skip("Websocket not enabled, skipping")
 	}
-	if k.WebsocketConn == nil {
-		k.Websocket.Connect()
-	}
 	var obUpdates []string
 	obpartial := `[0,{"as":[["5541.30000","2.50700000","0"]],"bs":[["5541.20000","1.52900000","0"]]}]`
 	for i := 1; i < orderbookBufferLimit+2; i++ {
@@ -704,9 +701,6 @@ func TestOrderBookOutOfOrder(t *testing.T) {
 	if !k.Websocket.IsEnabled() {
 		t.Skip("Websocket not enabled, skipping")
 	}
-	if k.WebsocketConn == nil {
-		k.Websocket.Connect()
-	}
 	obpartial := `[0,{"as":[["5541.30000","2.50700000","0"]],"bs":[["5541.20000","1.52900000","5"]]}]`
 	obupdate1 := `[0,{"a":[["5541.30000","0.00000000","1"]],"b":[["5541.30000","0.00000000","3"]]}]`
 	obupdate2 := `[0,{"a":[["5541.30000","2.50700000","2"]],"b":[["5541.30000","0.00000000","1"]]}]`
--- a/exchanges/okcoin/okcoin_test.go
+++ b/exchanges/okcoin/okcoin_test.go
@@ -1,10 +1,13 @@
 package okcoin

 import (
+	"net/http"
 	"strings"
+	"sync"
 	"testing"
 	"time"

+	"github.com/gorilla/websocket"
 	"github.com/thrasher-/gocryptotrader/common"
 	"github.com/thrasher-/gocryptotrader/config"
 	"github.com/thrasher-/gocryptotrader/currency"
@@ -92,30 +95,6 @@ func testStandardErrorHandling(t *testing.T, err error) {
 	}
 }

-// setupWSConnection Connect to WS, but pass back error so test can handle it if needed
-func setupWSConnection() error {
-	o.Enabled = true
-	err := o.WebsocketSetup(o.WsConnect,
-		nil,
-		nil,
-		o.Name,
-		true,
-		o.Verbose,
-		o.WebsocketURL,
-		o.WebsocketURL)
-	o.Websocket.DataHandler = make(chan interface{}, 500)
-	if err != nil {
-		return err
-	}
-	o.Websocket.SetWsStatusAndConnection(true)
-	return nil
-}
-
-// disconnectFromWS disconnect to WS, but pass back error so test can handle it if needed
-func disconnectFromWS() error {
-	return o.Websocket.Shutdown()
-}
-
 // TestGetAccountCurrencies API endpoint test
 func TestGetAccountCurrencies(t *testing.T) {
 	TestSetDefaults(t)
@@ -816,88 +795,52 @@ func TestGetMarginTransactionDetails(t *testing.T) {

 // Websocket tests ----------------------------------------------------------------------------------------------

-// TestWsLogin API endpoint test
-func TestWsLogin(t *testing.T) {
-	TestSetRealOrderDefaults(t)
-	if !websocketEnabled {
-		t.Skip("Websocket not enabled, skipping")
-	}
-	if !o.Websocket.IsConnecting() || !o.Websocket.IsConnected() {
-		o.Websocket.Connect()
-	}
-	err := o.WsLogin()
-	if err != nil {
-		t.Error(err)
-	}
-	var errorReceived bool
-	for i := 0; i < 5; i++ {
-		response := <-o.Websocket.DataHandler
-		if err, ok := response.(error); ok && err != nil {
-			errorReceived = true
-		}
-	}
-	if errorReceived {
-		t.Error("Expecting no errors")
-	}
-}
-
-// TestSubscribeToChannel API endpoint test
-func TestSubscribeToChannel(t *testing.T) {
-	TestSetDefaults(t)
-	if !websocketEnabled {
-		t.Skip("Websocket not enabled, skipping")
-	}
-	if !o.Websocket.IsConnecting() || !o.Websocket.IsConnected() {
-		o.Websocket.Connect()
-	}
-	subscription := exchange.WebsocketChannelSubscription{
-		Channel:  "spot/depth",
-		Currency: currency.NewPairDelimiter("LTC-BTC", "-"),
-	}
-
-	o.Subscribe(subscription)
-	var errorReceived bool
-	for i := 0; i < 5; i++ {
-		response := <-o.Websocket.DataHandler
-		if err, ok := response.(error); ok && err != nil {
-			t.Log(response)
-			if strings.Contains(response.(error).Error(), subscription.Channel) {
-				errorReceived = true
-			}
-		}
-	}
-	if errorReceived {
-		t.Error("Expecting subscription to channel")
-	}
-}
-
-// TestSubscribeToNonExistantChannel Logic test
+// TestSendWsMessages Logic test
 // Attempts to subscribe to a channel that doesn't exist
-// Then captures the error response
-func TestSubscribeToNonExistantChannel(t *testing.T) {
+// Will log in if credentials are present
+func TestSendWsMessages(t *testing.T) {
 	TestSetDefaults(t)
 	if !websocketEnabled {
 		t.Skip("Websocket not enabled, skipping")
 	}
-	if !o.Websocket.IsConnecting() || !o.Websocket.IsConnected() {
-		o.Websocket.Connect()
+	var dialer websocket.Dialer
+	var err error
+	var ok bool
+	o.Websocket.TrafficAlert = make(chan struct{}, 99)
+	o.WebsocketConn, _, err = dialer.Dial(o.Websocket.GetWebsocketURL(),
+		http.Header{})
+	if err != nil {
+		t.Fatalf("%s Unable to connect to Websocket. Error: %s",
+			o.Name,
+			err)
 	}
+	defer o.WebsocketConn.Close()
+	wg := sync.WaitGroup{}
+	wg.Add(1)
+	go o.WsHandleData(&wg)
+	wg.Wait()
+
 	subscription := exchange.WebsocketChannelSubscription{
 		Channel: "badChannel",
 	}
 	o.Subscribe(subscription)
-	var errorReceived bool
-	for i := 0; i < 5; i++ {
-		response := <-o.Websocket.DataHandler
-		if err, ok := response.(error); ok && err != nil {
-			t.Log(response)
-			if strings.Contains(response.(error).Error(), subscription.Channel) {
-				errorReceived = true
-			}
+	response := <-o.Websocket.DataHandler
+	if err, ok = response.(error); ok && err != nil {
+		if !strings.Contains(response.(error).Error(), subscription.Channel) {
+			t.Error("Expecting OKEX error - 30040 message: Channel badChannel doesn't exist")
 		}
 	}
-	if !errorReceived {
-		t.Error("Expecting OKEX error - 30040 message: Channel badChannel doesn't exist")
+
+	if !areTestAPIKeysSet() {
+		return
+	}
+	err = o.WsLogin()
+	if err != nil {
+		t.Error(err)
+	}
+	response = <-o.Websocket.DataHandler
+	if err, ok := response.(error); ok && err != nil {
+		t.Error(err)
 	}
 }

--- a/exchanges/okex/okex_test.go
+++ b/exchanges/okex/okex_test.go
@@ -2,9 +2,13 @@ package okex

 import (
 	"fmt"
+	"net/http"
+	"strings"
+	"sync"
 	"testing"
 	"time"

+	"github.com/gorilla/websocket"
 	"github.com/thrasher-/gocryptotrader/common"
 	"github.com/thrasher-/gocryptotrader/config"
 	"github.com/thrasher-/gocryptotrader/currency"
@@ -92,26 +96,6 @@ func testStandardErrorHandling(t *testing.T, err error) {
 	}
 }

-// setupWSConnection Connect to WS, but pass back error so test can handle it if needed
-func setupWSConnection() error {
-	if !o.Websocket.IsEnabled() {
-		err := o.WebsocketSetup(o.WsConnect,
-			o.Subscribe,
-			o.Unsubscribe,
-			o.Name,
-			true,
-			o.Verbose,
-			o.WebsocketURL,
-			o.WebsocketURL)
-		o.Websocket.DataHandler = make(chan interface{}, 500)
-		if err != nil {
-			return err
-		}
-		o.Websocket.SetWsStatusAndConnection(true)
-	}
-	return nil
-}
-
 // TestGetAccountCurrencies API endpoint test
 func TestGetAccountCurrencies(t *testing.T) {
 	TestSetDefaults(t)
@@ -1576,29 +1560,52 @@ func TestGetETTSettlementPriceHistory(t *testing.T) {

 // Websocket tests ----------------------------------------------------------------------------------------------

-// TestWsLogin API endpoint test
-func TestWsLogin(t *testing.T) {
-	TestSetRealOrderDefaults(t)
+// TestSendWsMessages Logic test
+// Attempts to subscribe to a channel that doesn't exist
+// Will log in if credentials are present
+func TestSendWsMessages(t *testing.T) {
+	TestSetDefaults(t)
 	if !websocketEnabled {
 		t.Skip("Websocket not enabled, skipping")
 	}
-	if !o.Websocket.IsConnecting() || !o.Websocket.IsConnected() {
-		o.Websocket.Connect()
+	var dialer websocket.Dialer
+	var err error
+	var ok bool
+	o.Websocket.TrafficAlert = make(chan struct{}, 99)
+	o.WebsocketConn, _, err = dialer.Dial(o.Websocket.GetWebsocketURL(),
+		http.Header{})
+	if err != nil {
+		t.Fatalf("%s Unable to connect to Websocket. Error: %s",
+			o.Name,
+			err)
 	}
-	err := o.WsLogin()
+	defer o.WebsocketConn.Close()
+	wg := sync.WaitGroup{}
+	wg.Add(1)
+	go o.WsHandleData(&wg)
+	wg.Wait()
+
+	subscription := exchange.WebsocketChannelSubscription{
+		Channel: "badChannel",
+	}
+	o.Subscribe(subscription)
+	response := <-o.Websocket.DataHandler
+	if err, ok = response.(error); ok && err != nil {
+		if !strings.Contains(response.(error).Error(), subscription.Channel) {
+			t.Error("Expecting OKEX error - 30040 message: Channel badChannel doesn't exist")
+		}
+	}
+
+	if !areTestAPIKeysSet() {
+		return
+	}
+	err = o.WsLogin()
 	if err != nil {
 		t.Error(err)
 	}
-	var errorReceived bool
-	for i := 0; i < 5; i++ {
-		response := <-o.Websocket.DataHandler
-		if err, ok := response.(error); ok && err != nil {
-			t.Log(response)
-			errorReceived = true
-		}
-	}
-	if errorReceived {
-		t.Error("Expecting no errors")
+	response = <-o.Websocket.DataHandler
+	if err, ok := response.(error); ok && err != nil {
+		t.Error(err)
 	}
 }