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

     1  package writer
     2  
     3  import (
     4  	"bytes"
     5  	"net"
     6  	"path/filepath"
     7  	"sync"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/Jeffail/benthos/v3/lib/log"
    12  	"github.com/Jeffail/benthos/v3/lib/message"
    13  	"github.com/Jeffail/benthos/v3/lib/metrics"
    14  )
    15  
    16  func TestSocketBasic(t *testing.T) {
    17  	tmpDir := t.TempDir()
    18  
    19  	ln, err := net.Listen("unix", filepath.Join(tmpDir, "benthos.sock"))
    20  	if err != nil {
    21  		t.Fatalf("failed to listen on address: %v", err)
    22  	}
    23  	defer ln.Close()
    24  
    25  	conf := NewSocketConfig()
    26  	conf.Network = ln.Addr().Network()
    27  	conf.Address = ln.Addr().String()
    28  
    29  	wtr, err := NewSocket(conf, nil, log.Noop(), metrics.Noop())
    30  	if err != nil {
    31  		t.Fatal(err)
    32  	}
    33  
    34  	defer func() {
    35  		if err := wtr.WaitForClose(time.Second); err != nil {
    36  			t.Error(err)
    37  		}
    38  	}()
    39  
    40  	go func() {
    41  		if cerr := wtr.Connect(); cerr != nil {
    42  			t.Error(cerr)
    43  		}
    44  	}()
    45  
    46  	conn, err := ln.Accept()
    47  	if err != nil {
    48  		t.Fatal(err)
    49  	}
    50  
    51  	var buf bytes.Buffer
    52  
    53  	wg := sync.WaitGroup{}
    54  	wg.Add(1)
    55  	go func() {
    56  		conn.SetReadDeadline(time.Now().Add(time.Second * 5))
    57  		buf.ReadFrom(conn)
    58  		wg.Done()
    59  	}()
    60  
    61  	if err = wtr.Write(message.New([][]byte{[]byte("foo")})); err != nil {
    62  		t.Error(err)
    63  	}
    64  	if err = wtr.Write(message.New([][]byte{[]byte("bar\n")})); err != nil {
    65  		t.Error(err)
    66  	}
    67  	if err = wtr.Write(message.New([][]byte{[]byte("baz")})); err != nil {
    68  		t.Error(err)
    69  	}
    70  	wtr.CloseAsync()
    71  	wg.Wait()
    72  
    73  	exp := "foo\nbar\nbaz\n"
    74  	if act := buf.String(); exp != act {
    75  		t.Errorf("Wrong result: %v != %v", act, exp)
    76  	}
    77  
    78  	conn.Close()
    79  }
    80  
    81  func TestSocketMultipart(t *testing.T) {
    82  	tmpDir := t.TempDir()
    83  
    84  	ln, err := net.Listen("unix", filepath.Join(tmpDir, "benthos.sock"))
    85  	if err != nil {
    86  		t.Fatalf("failed to listen on address: %v", err)
    87  	}
    88  	defer ln.Close()
    89  
    90  	conf := NewSocketConfig()
    91  	conf.Network = ln.Addr().Network()
    92  	conf.Address = ln.Addr().String()
    93  
    94  	wtr, err := NewSocket(conf, nil, log.Noop(), metrics.Noop())
    95  	if err != nil {
    96  		t.Fatal(err)
    97  	}
    98  
    99  	defer func() {
   100  		if err := wtr.WaitForClose(time.Second); err != nil {
   101  			t.Error(err)
   102  		}
   103  	}()
   104  
   105  	go func() {
   106  		if cerr := wtr.Connect(); cerr != nil {
   107  			t.Error(cerr)
   108  		}
   109  	}()
   110  
   111  	conn, err := ln.Accept()
   112  	if err != nil {
   113  		t.Fatal(err)
   114  	}
   115  
   116  	var buf bytes.Buffer
   117  
   118  	wg := sync.WaitGroup{}
   119  	wg.Add(1)
   120  	go func() {
   121  		conn.SetReadDeadline(time.Now().Add(time.Second * 5))
   122  		buf.ReadFrom(conn)
   123  		wg.Done()
   124  	}()
   125  
   126  	if err = wtr.Write(message.New([][]byte{[]byte("foo"), []byte("bar"), []byte("baz")})); err != nil {
   127  		t.Error(err)
   128  	}
   129  	if err = wtr.Write(message.New([][]byte{[]byte("qux")})); err != nil {
   130  		t.Error(err)
   131  	}
   132  	wtr.CloseAsync()
   133  	wg.Wait()
   134  
   135  	exp := "foo\nbar\nbaz\n\nqux\n"
   136  	if act := buf.String(); exp != act {
   137  		t.Errorf("Wrong result: %v != %v", act, exp)
   138  	}
   139  
   140  	conn.Close()
   141  }
   142  
   143  type wrapPacketConn struct {
   144  	r net.PacketConn
   145  }
   146  
   147  func (w *wrapPacketConn) Read(p []byte) (n int, err error) {
   148  	n, _, err = w.r.ReadFrom(p)
   149  	return
   150  }
   151  
   152  func TestUDPSocketBasic(t *testing.T) {
   153  	conn, err := net.ListenPacket("udp", "127.0.0.1:0")
   154  	if err != nil {
   155  		if conn, err = net.ListenPacket("tcp6", "[::1]:0"); err != nil {
   156  			t.Fatalf("failed to listen on a port: %v", err)
   157  		}
   158  	}
   159  	defer conn.Close()
   160  
   161  	conf := NewSocketConfig()
   162  	conf.Network = "udp"
   163  	conf.Address = conn.LocalAddr().String()
   164  
   165  	wtr, err := NewSocket(conf, nil, log.Noop(), metrics.Noop())
   166  	if err != nil {
   167  		t.Fatal(err)
   168  	}
   169  
   170  	defer func() {
   171  		if err := wtr.WaitForClose(time.Second); err != nil {
   172  			t.Error(err)
   173  		}
   174  	}()
   175  
   176  	if cerr := wtr.Connect(); cerr != nil {
   177  		t.Fatal(cerr)
   178  	}
   179  
   180  	var buf bytes.Buffer
   181  
   182  	wg := sync.WaitGroup{}
   183  	wg.Add(1)
   184  	go func() {
   185  		conn.SetReadDeadline(time.Now().Add(time.Second * 5))
   186  		buf.ReadFrom(&wrapPacketConn{r: conn})
   187  		wg.Done()
   188  	}()
   189  
   190  	if err = wtr.Write(message.New([][]byte{[]byte("foo")})); err != nil {
   191  		t.Error(err)
   192  	}
   193  	if err = wtr.Write(message.New([][]byte{[]byte("bar\n")})); err != nil {
   194  		t.Error(err)
   195  	}
   196  	if err = wtr.Write(message.New([][]byte{[]byte("baz")})); err != nil {
   197  		t.Error(err)
   198  	}
   199  	wtr.CloseAsync()
   200  	wg.Wait()
   201  
   202  	exp := "foo\nbar\nbaz\n"
   203  	if act := buf.String(); exp != act {
   204  		t.Errorf("Wrong result: %v != %v", act, exp)
   205  	}
   206  
   207  	conn.Close()
   208  }
   209  
   210  func TestUDPSocketMultipart(t *testing.T) {
   211  	conn, err := net.ListenPacket("udp", "127.0.0.1:0")
   212  	if err != nil {
   213  		if conn, err = net.ListenPacket("tcp6", "[::1]:0"); err != nil {
   214  			t.Fatalf("failed to listen on a port: %v", err)
   215  		}
   216  	}
   217  	defer conn.Close()
   218  
   219  	conf := NewSocketConfig()
   220  	conf.Network = "udp"
   221  	conf.Address = conn.LocalAddr().String()
   222  
   223  	wtr, err := NewSocket(conf, nil, log.Noop(), metrics.Noop())
   224  	if err != nil {
   225  		t.Fatal(err)
   226  	}
   227  
   228  	defer func() {
   229  		if err := wtr.WaitForClose(time.Second); err != nil {
   230  			t.Error(err)
   231  		}
   232  	}()
   233  
   234  	if cerr := wtr.Connect(); cerr != nil {
   235  		t.Fatal(cerr)
   236  	}
   237  
   238  	var buf bytes.Buffer
   239  
   240  	wg := sync.WaitGroup{}
   241  	wg.Add(1)
   242  	go func() {
   243  		conn.SetReadDeadline(time.Now().Add(time.Second * 5))
   244  		buf.ReadFrom(&wrapPacketConn{r: conn})
   245  		wg.Done()
   246  	}()
   247  
   248  	if err = wtr.Write(message.New([][]byte{[]byte("foo"), []byte("bar"), []byte("baz")})); err != nil {
   249  		t.Error(err)
   250  	}
   251  	if err = wtr.Write(message.New([][]byte{[]byte("qux")})); err != nil {
   252  		t.Error(err)
   253  	}
   254  	wtr.CloseAsync()
   255  	wg.Wait()
   256  
   257  	exp := "foo\nbar\nbaz\n\nqux\n"
   258  	if act := buf.String(); exp != act {
   259  		t.Errorf("Wrong result: %v != %v", act, exp)
   260  	}
   261  
   262  	conn.Close()
   263  }
   264  
   265  func TestTCPSocketBasic(t *testing.T) {
   266  	ln, err := net.Listen("tcp", "127.0.0.1:0")
   267  	if err != nil {
   268  		if ln, err = net.Listen("tcp6", "[::1]:0"); err != nil {
   269  			t.Fatalf("failed to listen on a port: %v", err)
   270  		}
   271  	}
   272  	defer ln.Close()
   273  
   274  	conf := NewSocketConfig()
   275  	conf.Network = "tcp"
   276  	conf.Address = ln.Addr().String()
   277  
   278  	wtr, err := NewSocket(conf, nil, log.Noop(), metrics.Noop())
   279  	if err != nil {
   280  		t.Fatal(err)
   281  	}
   282  
   283  	defer func() {
   284  		if err := wtr.WaitForClose(time.Second); err != nil {
   285  			t.Error(err)
   286  		}
   287  	}()
   288  
   289  	go func() {
   290  		if cerr := wtr.Connect(); cerr != nil {
   291  			t.Error(cerr)
   292  		}
   293  	}()
   294  
   295  	conn, err := ln.Accept()
   296  	if err != nil {
   297  		t.Fatal(err)
   298  	}
   299  
   300  	var buf bytes.Buffer
   301  
   302  	wg := sync.WaitGroup{}
   303  	wg.Add(1)
   304  	go func() {
   305  		conn.SetReadDeadline(time.Now().Add(time.Second * 5))
   306  		buf.ReadFrom(conn)
   307  		wg.Done()
   308  	}()
   309  
   310  	if err = wtr.Write(message.New([][]byte{[]byte("foo")})); err != nil {
   311  		t.Error(err)
   312  	}
   313  	if err = wtr.Write(message.New([][]byte{[]byte("bar\n")})); err != nil {
   314  		t.Error(err)
   315  	}
   316  	if err = wtr.Write(message.New([][]byte{[]byte("baz")})); err != nil {
   317  		t.Error(err)
   318  	}
   319  	wtr.CloseAsync()
   320  	wg.Wait()
   321  
   322  	exp := "foo\nbar\nbaz\n"
   323  	if act := buf.String(); exp != act {
   324  		t.Errorf("Wrong result: %v != %v", act, exp)
   325  	}
   326  
   327  	conn.Close()
   328  }
   329  
   330  func TestTCPSocketMultipart(t *testing.T) {
   331  	ln, err := net.Listen("tcp", "127.0.0.1:0")
   332  	if err != nil {
   333  		if ln, err = net.Listen("tcp6", "[::1]:0"); err != nil {
   334  			t.Fatalf("failed to listen on a port: %v", err)
   335  		}
   336  	}
   337  	defer ln.Close()
   338  
   339  	conf := NewSocketConfig()
   340  	conf.Network = "tcp"
   341  	conf.Address = ln.Addr().String()
   342  
   343  	wtr, err := NewSocket(conf, nil, log.Noop(), metrics.Noop())
   344  	if err != nil {
   345  		t.Fatal(err)
   346  	}
   347  
   348  	defer func() {
   349  		if err := wtr.WaitForClose(time.Second); err != nil {
   350  			t.Error(err)
   351  		}
   352  	}()
   353  
   354  	go func() {
   355  		if cerr := wtr.Connect(); cerr != nil {
   356  			t.Error(cerr)
   357  		}
   358  	}()
   359  
   360  	conn, err := ln.Accept()
   361  	if err != nil {
   362  		t.Fatal(err)
   363  	}
   364  
   365  	var buf bytes.Buffer
   366  
   367  	wg := sync.WaitGroup{}
   368  	wg.Add(1)
   369  	go func() {
   370  		conn.SetReadDeadline(time.Now().Add(time.Second * 5))
   371  		buf.ReadFrom(conn)
   372  		wg.Done()
   373  	}()
   374  
   375  	if err = wtr.Write(message.New([][]byte{[]byte("foo"), []byte("bar"), []byte("baz")})); err != nil {
   376  		t.Error(err)
   377  	}
   378  	if err = wtr.Write(message.New([][]byte{[]byte("qux")})); err != nil {
   379  		t.Error(err)
   380  	}
   381  	wtr.CloseAsync()
   382  	wg.Wait()
   383  
   384  	exp := "foo\nbar\nbaz\n\nqux\n"
   385  	if act := buf.String(); exp != act {
   386  		t.Errorf("Wrong result: %v != %v", act, exp)
   387  	}
   388  
   389  	conn.Close()
   390  }
   391  
   392  func TestSocketCustomDelimeter(t *testing.T) {
   393  	tmpDir := t.TempDir()
   394  
   395  	ln, err := net.Listen("unix", filepath.Join(tmpDir, "benthos.sock"))
   396  	if err != nil {
   397  		t.Fatalf("failed to listen on address: %v", err)
   398  	}
   399  	defer ln.Close()
   400  
   401  	conf := NewSocketConfig()
   402  	conf.Network = ln.Addr().Network()
   403  	conf.Address = ln.Addr().String()
   404  	conf.Codec = "delim:\t"
   405  
   406  	wtr, err := NewSocket(conf, nil, log.Noop(), metrics.Noop())
   407  	if err != nil {
   408  		t.Fatal(err)
   409  	}
   410  
   411  	defer func() {
   412  		if err := wtr.WaitForClose(time.Second); err != nil {
   413  			t.Error(err)
   414  		}
   415  	}()
   416  
   417  	go func() {
   418  		if cerr := wtr.Connect(); cerr != nil {
   419  			t.Error(cerr)
   420  		}
   421  	}()
   422  
   423  	conn, err := ln.Accept()
   424  	if err != nil {
   425  		t.Fatal(err)
   426  	}
   427  
   428  	var buf bytes.Buffer
   429  
   430  	wg := sync.WaitGroup{}
   431  	wg.Add(1)
   432  	go func() {
   433  		conn.SetReadDeadline(time.Now().Add(time.Second * 5))
   434  		buf.ReadFrom(conn)
   435  		wg.Done()
   436  	}()
   437  
   438  	if err = wtr.Write(message.New([][]byte{[]byte("foo")})); err != nil {
   439  		t.Error(err)
   440  	}
   441  	if err = wtr.Write(message.New([][]byte{[]byte("bar\n")})); err != nil {
   442  		t.Error(err)
   443  	}
   444  	if err = wtr.Write(message.New([][]byte{[]byte("baz\t")})); err != nil {
   445  		t.Error(err)
   446  	}
   447  	wtr.CloseAsync()
   448  	wg.Wait()
   449  
   450  	exp := "foo\tbar\n\tbaz\t"
   451  	if act := buf.String(); exp != act {
   452  		t.Errorf("Wrong result: %v != %v", act, exp)
   453  	}
   454  
   455  	conn.Close()
   456  }