mtg/mtglib/internal/telegram/telegram_internal_test.go

package telegram

import (
	"context"
	"errors"
	"io"
	"net"
	"strconv"
	"testing"

	"github.com/9seconds/mtg/v2/internal/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", false)
}

func (suite *TelegramTestSuite) TearDownTest() {
	suite.dialerMock.AssertExpectations(suite.T())
}

func (suite *TelegramTestSuite) TestUnknownDC() {
	testData := []int{
		-1,
		0,
		6,
		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)
			assert.False(t, suite.t.IsKnownDC(value))
		})
	}
}

func (suite *TelegramTestSuite) TestDialToCorrectIPs() {
	testData := map[int][]tgAddr{}

	for i := 1; i <= 5; i++ {
		testData[i] = []tgAddr{}
		testData[i] = append(testData[i], productionV4Addresses[i-1]...)
		testData[i] = append(testData[i], productionV6Addresses[i-1]...)
	}

	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))
			assert.True(t, suite.t.IsKnownDC(idx))
		})
	}
}

func (suite *TelegramTestSuite) TestDialPreferIPRange() {
	testData := map[string][]tgAddr{
		"prefer-ipv4": {testV4Addresses[0][0], testV6Addresses[0][0]},
		"prefer-ipv6": {testV6Addresses[0][0], testV4Addresses[0][0]},
		"only-ipv4":   {testV4Addresses[0][0]},
		"only-ipv6":   {testV6Addresses[0][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, true)
			_, err := tg.Dial(context.Background(), 1)

			assert.True(t, errors.Is(err, io.EOF))
		})
	}
}

func (suite *TelegramTestSuite) TestDialPreferIPPriority() {
	testData := map[string]tgAddr{
		"prefer-ipv4": productionV4Addresses[0][0],
		"prefer-ipv6": productionV6Addresses[0][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, false)

			res, err := tg.Dial(context.Background(), 1)
			assert.NoError(t, err)
			assert.Equal(t, conn, res)
		})
	}
}

func (suite *TelegramTestSuite) TestUnknownPreferIP() {
	_, err := New(suite.dialerMock, "xxx", false)
	suite.Error(err)
}

func (suite *TelegramTestSuite) TestFallbackDC() {
	dcs := make([]int, 10)

	for i := 0; i < len(dcs); i++ {
		dcs[i] = suite.t.GetFallbackDC()
	}

	for _, v := range dcs {
		value := v

		suite.T().Run(strconv.Itoa(value), func(t *testing.T) {
			assert.True(t, suite.t.IsKnownDC(value))
		})
	}
}

func TestTelegram(t *testing.T) {
	t.Parallel()
	suite.Run(t, &TelegramTestSuite{})
}