mtg/network/v2/echo_server_test.go

package network_test

import (
	"context"
	"io"
	"net"
	"sync"

	"github.com/stretchr/testify/require"
	"github.com/stretchr/testify/suite"
)

type EchoServer struct {
	wg        sync.WaitGroup
	ctx       context.Context
	ctxCancel context.CancelFunc
	listener  net.Listener
}

func (e *EchoServer) Run() {
	e.wg.Go(func() {
		<-e.ctx.Done()
		e.listener.Close() //nolint: errcheck
	})

	e.wg.Go(func() {
		for {
			conn, err := e.listener.Accept()
			if err != nil {
				return
			}

			e.wg.Go(func() {
				<-e.ctx.Done()
				conn.Close() //nolint: errcheck
			})
			e.wg.Go(func() {
				e.process(conn)
			})
		}
	})
}

func (e *EchoServer) Stop() {
	e.ctxCancel()
	e.wg.Wait()
}

func (e *EchoServer) Addr() string {
	return e.listener.Addr().String()
}

func (e *EchoServer) process(conn io.ReadWriter) {
	buf := [4096]byte{}

	for {
		select {
		case <-e.ctx.Done():
			return
		default:
		}

		n, err := conn.Read(buf[:])
		if err != nil {
			return
		}

		select {
		case <-e.ctx.Done():
			return
		default:
		}

		if _, err = conn.Write(buf[:n]); err != nil {
			return
		}
	}
}

type EchoServerTestSuite struct {
	suite.Suite

	echoServer *EchoServer
}

func (suite *EchoServerTestSuite) SetupSuite() {
	ctx, cancel := context.WithCancel(context.Background())

	listener, err := net.Listen("tcp", "127.0.0.1:0")
	require.NoError(suite.T(), err)

	suite.echoServer = &EchoServer{
		ctx:       ctx,
		ctxCancel: cancel,
		listener:  listener,
	}
	suite.echoServer.Run()
}

func (suite *EchoServerTestSuite) TearDownSuite() {
	suite.echoServer.Stop()
}

func (suite *EchoServerTestSuite) EchoServerAddr() string {
	return suite.echoServer.Addr()
}