Change algorithm of TCP relaying

mtglib/internal/relay/relay_test.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"io"
 	"testing"
-	"time"
 
 	"github.com/9seconds/mtg/v2/internal/testlib"
 	"github.com/9seconds/mtg/v2/mtglib/internal/relay"
@@ -15,60 +14,40 @@ import (
 type RelayTestSuite struct {
 	suite.Suite
 
-	ctx       context.Context
-	ctxCancel context.CancelFunc
-	r         *relay.Relay
+	loggerMock       relay.Logger
+	ctx              context.Context
+	ctxCancel        context.CancelFunc
+	telegramConnMock *testlib.NetConnMock
+	clientConnMock   *testlib.NetConnMock
 }
 
 func (suite *RelayTestSuite) SetupTest() {
-	suite.ctx, suite.ctxCancel = context.WithCancel(context.Background())
-	suite.r = relay.AcquireRelay(suite.ctx, loggerMock{}, 4096, time.Second)
+	ctx, cancel := context.WithCancel(context.Background())
+	suite.ctx = ctx
+	suite.ctxCancel = cancel
+	suite.loggerMock = &loggerMock{}
+	suite.telegramConnMock = &testlib.NetConnMock{}
+	suite.clientConnMock = &testlib.NetConnMock{}
 }
 
 func (suite *RelayTestSuite) TearDownTest() {
 	suite.ctxCancel()
-	relay.ReleaseRelay(suite.r)
-	suite.r = nil
+	suite.telegramConnMock.AssertExpectations(suite.T())
+	suite.clientConnMock.AssertExpectations(suite.T())
 }
 
-func (suite *RelayTestSuite) TestCancelled() {
-	suite.ctxCancel()
+func (suite *RelayTestSuite) TestExit() {
+	suite.telegramConnMock.On("SetReadDeadline", mock.Anything).Return(nil)
+	suite.telegramConnMock.On("Close").Return(nil)
+	suite.telegramConnMock.On("Read", mock.Anything).Return(10, io.EOF).Once()
+	suite.telegramConnMock.On("Write", mock.Anything).Return(10, io.EOF).Maybe()
 
-	eastConn := &rwcMock{}
-	eastConn.Write([]byte{1, 2, 3, 4, 5}) // nolint: errcheck
+	suite.clientConnMock.On("Read", mock.Anything).Return(0, io.EOF).Once()
+	suite.clientConnMock.On("Write", mock.Anything).Return(10, io.EOF).Maybe()
+	suite.clientConnMock.On("Close").Return(nil)
 
-	westConn := &rwcMock{}
-	westConn.Write([]byte{100, 101, 102}) // nolint: errcheck
-
-	suite.Nil(suite.r.Process(eastConn, westConn))
-}
-
-func (suite *RelayTestSuite) TestCopyFine() {
-	eastConn := &rwcMock{}
-	eastConn.Write([]byte{1, 2, 3, 4, 5}) // nolint: errcheck
-
-	westConn := &rwcMock{}
-	westConn.Write([]byte{100, 101, 102}) // nolint: errcheck
-
-	// yes, this test is not good enough. but apparently, if it hangs,
-	// we can debug most of possible issues.
-	_ = suite.r.Process(eastConn, westConn)
-}
-
-func (suite *RelayTestSuite) TestTimeout() {
-	eastConn := &rwcMock{}
-	eastConn.Write([]byte{1, 2, 3, 4, 5}) // nolint: errcheck
-
-	westConn := &testlib.NetConnMock{}
-	westConn.On("Close").Return(nil)
-	westConn.On("Read", mock.Anything).Return(0, io.EOF).Run(func(_ mock.Arguments) {
-		time.Sleep(2 * time.Second)
-	})
-	westConn.On("Write", mock.Anything).Return(0, io.EOF).Run(func(_ mock.Arguments) {
-		time.Sleep(2 * time.Second)
-	})
-
-	suite.Error(suite.r.Process(eastConn, westConn))
+	relay.Relay(suite.ctx, suite.loggerMock, 1024,
+		suite.telegramConnMock, suite.clientConnMock)
 }
 
 func TestRelay(t *testing.T) {