github.com/Jeffail/benthos/v3@v3.65.0/lib/input/tcp_test.go (about)

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