Websockets: Tests for GetSubscription behaviour

exchanges/stream/websocket_test.go

@@ -53,6 +53,10 @@ type testResponse struct {
 	RequestID int64 `json:"reqid,omitempty"`
 }
 
+type testSubKey struct {
+	Mood string
+}
+
 var defaultSetup = &WebsocketSetup{
 	ExchangeConfig: &config.Exchange{
 		Features: &config.FeaturesConfig{
@@ -72,9 +76,9 @@ var defaultSetup = &WebsocketSetup{
 	GenerateSubscriptions: func() ([]ChannelSubscription, error) {
 		return []ChannelSubscription{
 			{Channel: "TestSub"},
-			{Channel: "TestSub2"},
-			{Channel: "TestSub3"},
-			{Channel: "TestSub4"},
+			{Channel: "TestSub2", Key: "purple"},
+			{Channel: "TestSub3", Key: testSubKey{"mauve"}},
+			{Channel: "TestSub4", Key: 42},
 		}, nil
 	},
 	Features: &protocol.Features{Subscribe: true, Unsubscribe: true},
@@ -513,7 +517,25 @@ func TestSubscribeUnsubscribe(t *testing.T) {
 	assert.NoError(t, err, "Generating test subscriptions should not error")
 	assert.ErrorIs(t, ws.UnsubscribeChannels(nil), errNoChannelsInArgs, "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(defaultChannelKey{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), errAlreadySubscribed, "Subscribe should error when already subscribed")
 	assert.ErrorIs(t, ws.SubscribeToChannels(nil), errNoChannelsInArgs, "Subscribe to nil should error")
 	assert.NoError(t, ws.UnsubscribeChannels(subs), "Unsubscribing should not error")