diff --git a/mtglib/internal/relay/init_test.go b/mtglib/internal/relay/init_test.go new file mode 100644 index 0000000..e19ecf4 --- /dev/null +++ b/mtglib/internal/relay/init_test.go @@ -0,0 +1,49 @@ +package relay_test + +import ( + "bytes" + "io" + "sync" +) + +type loggerMock struct{} + +func (l loggerMock) Printf(format string, args ...interface{}) {} + +type rwcMock struct { + bytes.Buffer + + closed bool + mutex sync.Mutex +} + +func (r *rwcMock) Read(p []byte) (int, error) { + r.mutex.Lock() + defer r.mutex.Unlock() + + if r.closed { + return 0, io.EOF + } + + return r.Buffer.Read(p) +} + +func (r *rwcMock) Write(p []byte) (int, error) { + r.mutex.Lock() + defer r.mutex.Unlock() + + if r.closed { + return 0, io.EOF + } + + return r.Buffer.Write(p) +} + +func (r *rwcMock) Close() error { + r.mutex.Lock() + defer r.mutex.Unlock() + + r.closed = true + + return nil +} diff --git a/mtglib/internal/relay/relay.go b/mtglib/internal/relay/relay.go index 2a6e1b5..2785b03 100644 --- a/mtglib/internal/relay/relay.go +++ b/mtglib/internal/relay/relay.go @@ -34,7 +34,7 @@ func (r *Relay) Process(eastConn, westConn io.ReadWriteCloser) error { westConn.Close() }() - go r.runObserver() + go r.runObserver(r.ctx) wg := &sync.WaitGroup{} wg.Add(2) // nolint: gomnd @@ -45,7 +45,12 @@ func (r *Relay) Process(eastConn, westConn io.ReadWriteCloser) error { wg.Wait() - return <-r.errorChannel + select { + case err := <-r.errorChannel: + return err + default: + return nil + } } func (r *Relay) transmit(src io.ReadCloser, dst io.WriteCloser, @@ -67,7 +72,7 @@ func (r *Relay) transmit(src io.ReadCloser, dst io.WriteCloser, } } -func (r *Relay) runObserver() { +func (r *Relay) runObserver(ctx context.Context) { ticker := time.NewTicker(time.Second) defer func() { @@ -80,7 +85,6 @@ func (r *Relay) runObserver() { }() lastTickAt := time.Now() - ctx := r.ctx for { select { diff --git a/mtglib/internal/relay/relay_test.go b/mtglib/internal/relay/relay_test.go new file mode 100644 index 0000000..b9d7987 --- /dev/null +++ b/mtglib/internal/relay/relay_test.go @@ -0,0 +1,56 @@ +package relay_test + +import ( + "context" + "testing" + "time" + + "github.com/9seconds/mtg/v2/mtglib/internal/relay" + "github.com/stretchr/testify/suite" +) + +type RelayTestSuite struct { + suite.Suite + + ctx context.Context + ctxCancel context.CancelFunc + r *relay.Relay +} + +func (suite *RelayTestSuite) SetupTest() { + suite.ctx, suite.ctxCancel = context.WithCancel(context.Background()) + suite.r = relay.AcquireRelay(suite.ctx, loggerMock{}, 4096, time.Second) +} + +func (suite *RelayTestSuite) TearDownTest() { + suite.ctxCancel() + relay.ReleaseRelay(suite.r) + suite.r = nil +} + +func (suite *RelayTestSuite) TestCancelled() { + suite.ctxCancel() + + eastConn := &rwcMock{} + eastConn.Write([]byte{1, 2, 3, 4, 5}) + + westConn := &rwcMock{} + westConn.Write([]byte{100, 101, 102}) + + suite.Nil(suite.r.Process(eastConn, westConn)) +} + +func (suite *RelayTestSuite) TestCopyFine() { + eastConn := &rwcMock{} + eastConn.Write([]byte{1, 2, 3, 4, 5}) + + westConn := &rwcMock{} + westConn.Write([]byte{100, 101, 102}) + + suite.NotNil(suite.r.Process(eastConn, westConn)) +} + +func TestRelay(t *testing.T) { + t.Parallel() + suite.Run(t, &RelayTestSuite{}) +}