github.com/observiq/carbon@v0.9.11-0.20200820160507-1b872e368a5e/operator/builtin/input/tcp_test.go (about)

     1  package input
     2  
     3  import (
     4  	"net"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/observiq/carbon/entry"
     9  	"github.com/observiq/carbon/operator"
    10  	"github.com/observiq/carbon/testutil"
    11  	"github.com/stretchr/testify/mock"
    12  	"github.com/stretchr/testify/require"
    13  )
    14  
    15  func tcpInputTest(input []byte, expected []string) func(t *testing.T) {
    16  	return func(t *testing.T) {
    17  		cfg := NewTCPInputConfig("test_id")
    18  		cfg.ListenAddress = ":0"
    19  
    20  		newOperator, err := cfg.Build(testutil.NewBuildContext(t))
    21  		require.NoError(t, err)
    22  
    23  		mockOutput := testutil.Operator{}
    24  		tcpInput := newOperator.(*TCPInput)
    25  		tcpInput.InputOperator.OutputOperators = []operator.Operator{&mockOutput}
    26  
    27  		entryChan := make(chan *entry.Entry, 1)
    28  		mockOutput.On("Process", mock.Anything, mock.Anything).Run(func(args mock.Arguments) {
    29  			entryChan <- args.Get(1).(*entry.Entry)
    30  		}).Return(nil)
    31  
    32  		err = tcpInput.Start()
    33  		require.NoError(t, err)
    34  		defer tcpInput.Stop()
    35  
    36  		conn, err := net.Dial("tcp", tcpInput.listener.Addr().String())
    37  		require.NoError(t, err)
    38  		defer conn.Close()
    39  
    40  		_, err = conn.Write(input)
    41  		require.NoError(t, err)
    42  
    43  		for _, expectedMessage := range expected {
    44  			select {
    45  			case entry := <-entryChan:
    46  				require.Equal(t, expectedMessage, entry.Record)
    47  			case <-time.After(time.Second):
    48  				require.FailNow(t, "Timed out waiting for message to be written")
    49  			}
    50  		}
    51  
    52  		select {
    53  		case entry := <-entryChan:
    54  			require.FailNow(t, "Unexpected entry: %s", entry)
    55  		case <-time.After(100 * time.Millisecond):
    56  			return
    57  		}
    58  	}
    59  }
    60  
    61  func TestTcpInput(t *testing.T) {
    62  	t.Run("Simple", tcpInputTest([]byte("message\n"), []string{"message"}))
    63  	t.Run("CarriageReturn", tcpInputTest([]byte("message\r\n"), []string{"message"}))
    64  }
    65  
    66  func BenchmarkTcpInput(b *testing.B) {
    67  	cfg := NewTCPInputConfig("test_id")
    68  	cfg.ListenAddress = ":0"
    69  
    70  	newOperator, err := cfg.Build(testutil.NewBuildContext(b))
    71  	require.NoError(b, err)
    72  
    73  	fakeOutput := testutil.NewFakeOutput(b)
    74  	tcpInput := newOperator.(*TCPInput)
    75  	tcpInput.InputOperator.OutputOperators = []operator.Operator{fakeOutput}
    76  
    77  	err = tcpInput.Start()
    78  	require.NoError(b, err)
    79  
    80  	done := make(chan struct{})
    81  	go func() {
    82  		conn, err := net.Dial("tcp", tcpInput.listener.Addr().String())
    83  		require.NoError(b, err)
    84  		defer tcpInput.Stop()
    85  		defer conn.Close()
    86  		message := []byte("message\n")
    87  		for {
    88  			select {
    89  			case <-done:
    90  				return
    91  			default:
    92  				_, err := conn.Write(message)
    93  				require.NoError(b, err)
    94  			}
    95  		}
    96  	}()
    97  
    98  	for i := 0; i < b.N; i++ {
    99  		<-fakeOutput.Received
   100  	}
   101  
   102  	defer close(done)
   103  }