github.com/Jeffail/benthos/v3@v3.65.0/lib/output/writer/tcp_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 TestTCPBasic(t *testing.T) {
    16  	ln, err := net.Listen("tcp", "127.0.0.1:0")
    17  	if err != nil {
    18  		if ln, err = net.Listen("tcp6", "[::1]:0"); err != nil {
    19  			t.Fatalf("failed to listen on a port: %v", err)
    20  		}
    21  	}
    22  	defer ln.Close()
    23  
    24  	conf := NewTCPConfig()
    25  	conf.Address = ln.Addr().String()
    26  
    27  	wtr, err := NewTCP(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  	go func() {
    39  		if cerr := wtr.Connect(); cerr != nil {
    40  			t.Error(cerr)
    41  		}
    42  	}()
    43  
    44  	conn, err := ln.Accept()
    45  	if err != nil {
    46  		t.Fatal(err)
    47  	}
    48  
    49  	var buf bytes.Buffer
    50  
    51  	wg := sync.WaitGroup{}
    52  	wg.Add(1)
    53  	go func() {
    54  		conn.SetReadDeadline(time.Now().Add(time.Second * 5))
    55  		buf.ReadFrom(conn)
    56  		wg.Done()
    57  	}()
    58  
    59  	if err = wtr.Write(message.New([][]byte{[]byte("foo")})); err != nil {
    60  		t.Error(err)
    61  	}
    62  	if err = wtr.Write(message.New([][]byte{[]byte("bar\n")})); err != nil {
    63  		t.Error(err)
    64  	}
    65  	if err = wtr.Write(message.New([][]byte{[]byte("baz")})); err != nil {
    66  		t.Error(err)
    67  	}
    68  	wtr.CloseAsync()
    69  	wg.Wait()
    70  
    71  	exp := "foo\nbar\nbaz\n"
    72  	if act := buf.String(); exp != act {
    73  		t.Errorf("Wrong result: %v != %v", act, exp)
    74  	}
    75  
    76  	conn.Close()
    77  }
    78  
    79  func TestTCPMultipart(t *testing.T) {
    80  	ln, err := net.Listen("tcp", "127.0.0.1:0")
    81  	if err != nil {
    82  		if ln, err = net.Listen("tcp6", "[::1]:0"); err != nil {
    83  			t.Fatalf("failed to listen on a port: %v", err)
    84  		}
    85  	}
    86  	defer ln.Close()
    87  
    88  	conf := NewTCPConfig()
    89  	conf.Address = ln.Addr().String()
    90  
    91  	wtr, err := NewTCP(conf, nil, log.Noop(), metrics.Noop())
    92  	if err != nil {
    93  		t.Fatal(err)
    94  	}
    95  
    96  	defer func() {
    97  		if err := wtr.WaitForClose(time.Second); err != nil {
    98  			t.Error(err)
    99  		}
   100  	}()
   101  
   102  	go func() {
   103  		if cerr := wtr.Connect(); cerr != nil {
   104  			t.Error(cerr)
   105  		}
   106  	}()
   107  
   108  	conn, err := ln.Accept()
   109  	if err != nil {
   110  		t.Fatal(err)
   111  	}
   112  
   113  	var buf bytes.Buffer
   114  
   115  	wg := sync.WaitGroup{}
   116  	wg.Add(1)
   117  	go func() {
   118  		conn.SetReadDeadline(time.Now().Add(time.Second * 5))
   119  		buf.ReadFrom(conn)
   120  		wg.Done()
   121  	}()
   122  
   123  	if err = wtr.Write(message.New([][]byte{[]byte("foo"), []byte("bar"), []byte("baz")})); err != nil {
   124  		t.Error(err)
   125  	}
   126  	if err = wtr.Write(message.New([][]byte{[]byte("qux")})); err != nil {
   127  		t.Error(err)
   128  	}
   129  	wtr.CloseAsync()
   130  	wg.Wait()
   131  
   132  	exp := "foo\nbar\nbaz\n\nqux\n"
   133  	if act := buf.String(); exp != act {
   134  		t.Errorf("Wrong result: %v != %v", act, exp)
   135  	}
   136  
   137  	conn.Close()
   138  }