From 925a02dac38f794193360336afde039434a66ea6 Mon Sep 17 00:00:00 2001
From: 9seconds <nineseconds@yandex.ru>
Date: Mon, 22 Mar 2021 14:51:18 +0300
Subject: [PATCH] Test telegram dc dialer

---
 mtglib/internal/telegram/telegram.go          |   2 +-
 .../telegram/telegram_internal_test.go        | 140 ++++++++++++++++++
 2 files changed, 141 insertions(+), 1 deletion(-)
 create mode 100644 mtglib/internal/telegram/telegram_internal_test.go

diff --git a/mtglib/internal/telegram/telegram.go b/mtglib/internal/telegram/telegram.go
index c8e3a27..90f21ca 100644
--- a/mtglib/internal/telegram/telegram.go
+++ b/mtglib/internal/telegram/telegram.go
@@ -52,7 +52,7 @@ func (t Telegram) Dial(ctx context.Context, dc int) (net.Conn, error) {
 	return nil, fmt.Errorf("cannot dial to %d dc: %w", dc, err)
 }
 
-func NewTelegramDialer(dialer Dialer, ipPreference string) (*Telegram, error) {
+func New(dialer Dialer, ipPreference string) (*Telegram, error) {
 	var pref preferIP
 
 	switch strings.ToLower(ipPreference) {
diff --git a/mtglib/internal/telegram/telegram_internal_test.go b/mtglib/internal/telegram/telegram_internal_test.go
new file mode 100644
index 0000000..891eff3
--- /dev/null
+++ b/mtglib/internal/telegram/telegram_internal_test.go
@@ -0,0 +1,140 @@
+package telegram
+
+import (
+	"context"
+	"errors"
+	"io"
+	"net"
+	"strconv"
+	"testing"
+
+	"github.com/9seconds/mtg/v2/testlib"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/suite"
+)
+
+type TelegramTestSuite struct {
+	suite.Suite
+
+	dialerMock *testlib.MtglibNetworkMock
+	t          *Telegram
+}
+
+func (suite *TelegramTestSuite) SetupTest() {
+	suite.dialerMock = &testlib.MtglibNetworkMock{}
+	suite.t, _ = New(suite.dialerMock, "prefer-ipv4")
+}
+
+func (suite *TelegramTestSuite) TearDownTest() {
+	suite.dialerMock.AssertExpectations(suite.T())
+}
+
+func (suite *TelegramTestSuite) TestUnknownDC() {
+	testData := []int{
+		-1,
+		5,
+		100,
+	}
+
+	for _, v := range testData {
+		value := v
+
+		suite.T().Run(strconv.Itoa(value), func(t *testing.T) {
+			_, err := suite.t.Dial(context.Background(), value)
+			assert.Error(t, err)
+		})
+	}
+}
+
+func (suite *TelegramTestSuite) TestDialToCorrectIPs() {
+	testData := map[int][]tgAddr{}
+
+	for i := 0; i < 5; i++ {
+		testData[i] = []tgAddr{}
+		testData[i] = append(testData[i], v4Addresses[i]...)
+		testData[i] = append(testData[i], v6Addresses[i])
+	}
+
+	for i, v := range testData {
+		idx := i
+		addresses := v
+
+		suite.T().Run(strconv.Itoa(idx), func(t *testing.T) {
+			for _, addr := range addresses {
+				suite.dialerMock.
+					On("DialContext", mock.Anything, addr.network, addr.address).
+					Once().
+					Return((*net.TCPConn)(nil), io.EOF)
+			}
+
+			_, err := suite.t.Dial(context.Background(), idx)
+			assert.True(t, errors.Is(err, io.EOF))
+		})
+	}
+}
+
+func (suite *TelegramTestSuite) TestDialPreferIPRange() {
+	testData := map[string][]tgAddr{
+		"prefer-ipv4": {v4Addresses[0][0], v6Addresses[0]},
+		"prefer-ipv6": {v6Addresses[0], v4Addresses[0][0]},
+		"only-ipv4":   {v4Addresses[0][0]},
+		"only-ipv6":   {v6Addresses[0]},
+	}
+
+	for k, v := range testData {
+		name := k
+		addresses := v
+
+		suite.T().Run(name, func(t *testing.T) {
+			for _, addr := range addresses {
+				suite.dialerMock.
+					On("DialContext", mock.Anything, addr.network, addr.address).
+					Once().
+					Return((*net.TCPConn)(nil), io.EOF)
+			}
+
+			tg, _ := New(suite.dialerMock, name)
+			_, err := tg.Dial(context.Background(), 0)
+
+			assert.True(t, errors.Is(err, io.EOF))
+		})
+	}
+}
+
+func (suite *TelegramTestSuite) TestDialPreferIPPriority() {
+	testData := map[string]tgAddr{
+		"prefer-ipv4": v4Addresses[0][0],
+		"prefer-ipv6": v6Addresses[0],
+	}
+
+	for k, v := range testData {
+		name := k
+		addr := v
+
+		suite.T().Run(name, func(t *testing.T) {
+			conn := &net.TCPConn{}
+
+			suite.dialerMock.
+				On("DialContext", mock.Anything, addr.network, addr.address).
+				Once().
+				Return(conn, nil)
+
+			tg, _ := New(suite.dialerMock, name)
+
+			res, err := tg.Dial(context.Background(), 0)
+			assert.NoError(t, err)
+			assert.Equal(t, conn, res)
+		})
+	}
+}
+
+func (suite *TelegramTestSuite) TestUnknownPreferIP() {
+	_, err := New(suite.dialerMock, "xxx")
+	suite.Error(err)
+}
+
+func TestTelegram(t *testing.T) {
+	t.Parallel()
+	suite.Run(t, &TelegramTestSuite{})
+}