github.com/observiq/carbon@v0.9.11-0.20200820160507-1b872e368a5e/operator/builtin/input/udp_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 udpInputTest(input []byte, expected []string) func(t *testing.T) {
    16  	return func(t *testing.T) {
    17  		cfg := NewUDPInputConfig("test_input")
    18  		cfg.ListenAddress = ":0"
    19  
    20  		newOperator, err := cfg.Build(testutil.NewBuildContext(t))
    21  		require.NoError(t, err)
    22  
    23  		mockOutput := testutil.Operator{}
    24  		udpInput, ok := newOperator.(*UDPInput)
    25  		require.True(t, ok)
    26  
    27  		udpInput.InputOperator.OutputOperators = []operator.Operator{&mockOutput}
    28  
    29  		entryChan := make(chan *entry.Entry, 1)
    30  		mockOutput.On("Process", mock.Anything, mock.Anything).Run(func(args mock.Arguments) {
    31  			entryChan <- args.Get(1).(*entry.Entry)
    32  		}).Return(nil)
    33  
    34  		err = udpInput.Start()
    35  		require.NoError(t, err)
    36  		defer udpInput.Stop()
    37  
    38  		conn, err := net.Dial("udp", udpInput.connection.LocalAddr().String())
    39  		require.NoError(t, err)
    40  		defer conn.Close()
    41  
    42  		_, err = conn.Write(input)
    43  		require.NoError(t, err)
    44  
    45  		for _, expectedRecord := range expected {
    46  			select {
    47  			case entry := <-entryChan:
    48  				require.Equal(t, expectedRecord, entry.Record)
    49  			case <-time.After(time.Second):
    50  				require.FailNow(t, "Timed out waiting for message to be written")
    51  			}
    52  		}
    53  
    54  		select {
    55  		case entry := <-entryChan:
    56  			require.FailNow(t, "Unexpected entry: %s", entry)
    57  		case <-time.After(100 * time.Millisecond):
    58  			return
    59  		}
    60  	}
    61  }
    62  
    63  func TestUDPInput(t *testing.T) {
    64  	t.Run("Simple", udpInputTest([]byte("message1"), []string{"message1"}))
    65  	t.Run("TrailingNewlines", udpInputTest([]byte("message1\n"), []string{"message1"}))
    66  	t.Run("TrailingCRNewlines", udpInputTest([]byte("message1\r\n"), []string{"message1"}))
    67  	t.Run("NewlineInMessage", udpInputTest([]byte("message1\nmessage2\n"), []string{"message1\nmessage2"}))
    68  }
    69  
    70  func BenchmarkUdpInput(b *testing.B) {
    71  	cfg := NewUDPInputConfig("test_id")
    72  	cfg.ListenAddress = ":0"
    73  
    74  	newOperator, err := cfg.Build(testutil.NewBuildContext(b))
    75  	require.NoError(b, err)
    76  
    77  	fakeOutput := testutil.NewFakeOutput(b)
    78  	udpInput := newOperator.(*UDPInput)
    79  	udpInput.InputOperator.OutputOperators = []operator.Operator{fakeOutput}
    80  
    81  	err = udpInput.Start()
    82  	require.NoError(b, err)
    83  
    84  	done := make(chan struct{})
    85  	go func() {
    86  		conn, err := net.Dial("udp", udpInput.connection.LocalAddr().String())
    87  		require.NoError(b, err)
    88  		defer udpInput.Stop()
    89  		defer conn.Close()
    90  		message := []byte("message\n")
    91  		for {
    92  			select {
    93  			case <-done:
    94  				return
    95  			default:
    96  				_, err := conn.Write(message)
    97  				require.NoError(b, err)
    98  			}
    99  		}
   100  	}()
   101  
   102  	for i := 0; i < b.N; i++ {
   103  		<-fakeOutput.Received
   104  	}
   105  
   106  	defer close(done)
   107  }