github.com/Jeffail/benthos/v3@v3.65.0/lib/output/writer/udp_test.go (about)

     1  package writer
     2  
     3  import (
     4  	"bytes"
     5  	"net"
     6  	"sync"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/Jeffail/benthos/v3/lib/log"
    11  	"github.com/Jeffail/benthos/v3/lib/message"
    12  	"github.com/Jeffail/benthos/v3/lib/metrics"
    13  )
    14  
    15  func TestUDPBasic(t *testing.T) {
    16  	conn, err := net.ListenPacket("udp", "127.0.0.1:0")
    17  	if err != nil {
    18  		if conn, err = net.ListenPacket("tcp6", "[::1]:0"); err != nil {
    19  			t.Fatalf("failed to listen on a port: %v", err)
    20  		}
    21  	}
    22  	defer conn.Close()
    23  
    24  	conf := NewUDPConfig()
    25  	conf.Address = conn.LocalAddr().String()
    26  
    27  	wtr, err := NewUDP(conf, nil, log.Noop(), metrics.Noop())
    28  	if err != nil {
    29  		t.Fatal(err)
    30  	}
    31  
    32  	defer func() {
    33  		if err := wtr.WaitForClose(time.Second); err != nil {
    34  			t.Error(err)
    35  		}
    36  	}()
    37  
    38  	if cerr := wtr.Connect(); cerr != nil {
    39  		t.Fatal(cerr)
    40  	}
    41  
    42  	var buf bytes.Buffer
    43  
    44  	wg := sync.WaitGroup{}
    45  	wg.Add(1)
    46  	go func() {
    47  		conn.SetReadDeadline(time.Now().Add(time.Second * 5))
    48  		buf.ReadFrom(&wrapPacketConn{r: conn})
    49  		wg.Done()
    50  	}()
    51  
    52  	if err = wtr.Write(message.New([][]byte{[]byte("foo")})); err != nil {
    53  		t.Error(err)
    54  	}
    55  	if err = wtr.Write(message.New([][]byte{[]byte("bar\n")})); err != nil {
    56  		t.Error(err)
    57  	}
    58  	if err = wtr.Write(message.New([][]byte{[]byte("baz")})); err != nil {
    59  		t.Error(err)
    60  	}
    61  	wtr.CloseAsync()
    62  	wg.Wait()
    63  
    64  	exp := "foo\nbar\nbaz\n"
    65  	if act := buf.String(); exp != act {
    66  		t.Errorf("Wrong result: %v != %v", act, exp)
    67  	}
    68  
    69  	conn.Close()
    70  }
    71  
    72  func TestUDPMultipart(t *testing.T) {
    73  	conn, err := net.ListenPacket("udp", "127.0.0.1:0")
    74  	if err != nil {
    75  		if conn, err = net.ListenPacket("tcp6", "[::1]:0"); err != nil {
    76  			t.Fatalf("failed to listen on a port: %v", err)
    77  		}
    78  	}
    79  	defer conn.Close()
    80  
    81  	conf := NewUDPConfig()
    82  	conf.Address = conn.LocalAddr().String()
    83  
    84  	wtr, err := NewUDP(conf, nil, log.Noop(), metrics.Noop())
    85  	if err != nil {
    86  		t.Fatal(err)
    87  	}
    88  
    89  	defer func() {
    90  		if err := wtr.WaitForClose(time.Second); err != nil {
    91  			t.Error(err)
    92  		}
    93  	}()
    94  
    95  	if cerr := wtr.Connect(); cerr != nil {
    96  		t.Fatal(cerr)
    97  	}
    98  
    99  	var buf bytes.Buffer
   100  
   101  	wg := sync.WaitGroup{}
   102  	wg.Add(1)
   103  	go func() {
   104  		conn.SetReadDeadline(time.Now().Add(time.Second * 5))
   105  		buf.ReadFrom(&wrapPacketConn{r: conn})
   106  		wg.Done()
   107  	}()
   108  
   109  	if err = wtr.Write(message.New([][]byte{[]byte("foo"), []byte("bar"), []byte("baz")})); err != nil {
   110  		t.Error(err)
   111  	}
   112  	if err = wtr.Write(message.New([][]byte{[]byte("qux")})); err != nil {
   113  		t.Error(err)
   114  	}
   115  	wtr.CloseAsync()
   116  	wg.Wait()
   117  
   118  	exp := "foo\nbar\nbaz\n\nqux\n"
   119  	if act := buf.String(); exp != act {
   120  		t.Errorf("Wrong result: %v != %v", act, exp)
   121  	}
   122  
   123  	conn.Close()
   124  }