github.com/Jeffail/benthos/v3@v3.65.0/public/service/input_test.go (about)

     1  package service
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/Jeffail/benthos/v3/lib/message"
    10  	"github.com/Jeffail/benthos/v3/lib/response"
    11  	"github.com/Jeffail/benthos/v3/lib/types"
    12  	"github.com/stretchr/testify/assert"
    13  )
    14  
    15  type fnInput struct {
    16  	connect func() error
    17  	read    func() (*Message, AckFunc, error)
    18  	closed  bool
    19  }
    20  
    21  func (f *fnInput) Connect(ctx context.Context) error {
    22  	return f.connect()
    23  }
    24  
    25  func (f *fnInput) Read(ctx context.Context) (*Message, AckFunc, error) {
    26  	return f.read()
    27  }
    28  
    29  func (f *fnInput) Close(ctx context.Context) error {
    30  	f.closed = true
    31  	return nil
    32  }
    33  
    34  func TestInputAirGapShutdown(t *testing.T) {
    35  	i := &fnInput{}
    36  	agi := newAirGapReader(i)
    37  
    38  	err := agi.WaitForClose(time.Millisecond * 5)
    39  	assert.EqualError(t, err, "action timed out")
    40  	assert.False(t, i.closed)
    41  
    42  	agi.CloseAsync()
    43  	err = agi.WaitForClose(time.Millisecond * 5)
    44  	assert.NoError(t, err)
    45  	assert.True(t, i.closed)
    46  }
    47  
    48  func TestInputAirGapSad(t *testing.T) {
    49  	i := &fnInput{
    50  		connect: func() error {
    51  			return errors.New("bad connect")
    52  		},
    53  		read: func() (*Message, AckFunc, error) {
    54  			return nil, nil, errors.New("bad read")
    55  		},
    56  	}
    57  	agi := newAirGapReader(i)
    58  
    59  	err := agi.ConnectWithContext(context.Background())
    60  	assert.EqualError(t, err, "bad connect")
    61  
    62  	_, _, err = agi.ReadWithContext(context.Background())
    63  	assert.EqualError(t, err, "bad read")
    64  
    65  	i.read = func() (*Message, AckFunc, error) {
    66  		return nil, nil, ErrNotConnected
    67  	}
    68  
    69  	_, _, err = agi.ReadWithContext(context.Background())
    70  	assert.Equal(t, types.ErrNotConnected, err)
    71  
    72  	i.read = func() (*Message, AckFunc, error) {
    73  		return nil, nil, ErrEndOfInput
    74  	}
    75  
    76  	_, _, err = agi.ReadWithContext(context.Background())
    77  	assert.Equal(t, types.ErrTypeClosed, err)
    78  }
    79  
    80  func TestInputAirGapHappy(t *testing.T) {
    81  	var ackErr error
    82  	ackFn := func(ctx context.Context, err error) error {
    83  		ackErr = err
    84  		return nil
    85  	}
    86  	i := &fnInput{
    87  		connect: func() error {
    88  			return nil
    89  		},
    90  		read: func() (*Message, AckFunc, error) {
    91  			m := &Message{
    92  				part: message.NewPart([]byte("hello world")),
    93  			}
    94  			return m, ackFn, nil
    95  		},
    96  	}
    97  	agi := newAirGapReader(i)
    98  
    99  	err := agi.ConnectWithContext(context.Background())
   100  	assert.NoError(t, err)
   101  
   102  	outMsg, outAckFn, err := agi.ReadWithContext(context.Background())
   103  	assert.NoError(t, err)
   104  	assert.Equal(t, 1, outMsg.Len())
   105  	assert.Equal(t, "hello world", string(outMsg.Get(0).Get()))
   106  
   107  	assert.NoError(t, outAckFn(context.Background(), response.NewError(errors.New("foobar"))))
   108  	assert.EqualError(t, ackErr, "foobar")
   109  }
   110  
   111  type fnBatchInput struct {
   112  	connect func() error
   113  	read    func() (MessageBatch, AckFunc, error)
   114  	closed  bool
   115  }
   116  
   117  func (f *fnBatchInput) Connect(ctx context.Context) error {
   118  	return f.connect()
   119  }
   120  
   121  func (f *fnBatchInput) ReadBatch(ctx context.Context) (MessageBatch, AckFunc, error) {
   122  	return f.read()
   123  }
   124  
   125  func (f *fnBatchInput) Close(ctx context.Context) error {
   126  	f.closed = true
   127  	return nil
   128  }
   129  
   130  func TestBatchInputAirGapShutdown(t *testing.T) {
   131  	i := &fnBatchInput{}
   132  	agi := newAirGapBatchReader(i)
   133  
   134  	err := agi.WaitForClose(time.Millisecond * 5)
   135  	assert.EqualError(t, err, "action timed out")
   136  	assert.False(t, i.closed)
   137  
   138  	agi.CloseAsync()
   139  	err = agi.WaitForClose(time.Millisecond * 5)
   140  	assert.NoError(t, err)
   141  	assert.True(t, i.closed)
   142  }
   143  
   144  func TestBatchInputAirGapSad(t *testing.T) {
   145  	i := &fnBatchInput{
   146  		connect: func() error {
   147  			return errors.New("bad connect")
   148  		},
   149  		read: func() (MessageBatch, AckFunc, error) {
   150  			return nil, nil, errors.New("bad read")
   151  		},
   152  	}
   153  	agi := newAirGapBatchReader(i)
   154  
   155  	err := agi.ConnectWithContext(context.Background())
   156  	assert.EqualError(t, err, "bad connect")
   157  
   158  	_, _, err = agi.ReadWithContext(context.Background())
   159  	assert.EqualError(t, err, "bad read")
   160  
   161  	i.read = func() (MessageBatch, AckFunc, error) {
   162  		return nil, nil, ErrNotConnected
   163  	}
   164  
   165  	_, _, err = agi.ReadWithContext(context.Background())
   166  	assert.Equal(t, types.ErrNotConnected, err)
   167  
   168  	i.read = func() (MessageBatch, AckFunc, error) {
   169  		return nil, nil, ErrEndOfInput
   170  	}
   171  
   172  	_, _, err = agi.ReadWithContext(context.Background())
   173  	assert.Equal(t, types.ErrTypeClosed, err)
   174  }
   175  
   176  func TestBatchInputAirGapHappy(t *testing.T) {
   177  	var ackErr error
   178  	ackFn := func(ctx context.Context, err error) error {
   179  		ackErr = err
   180  		return nil
   181  	}
   182  	i := &fnBatchInput{
   183  		connect: func() error {
   184  			return nil
   185  		},
   186  		read: func() (MessageBatch, AckFunc, error) {
   187  			m := MessageBatch{
   188  				NewMessage([]byte("hello world")),
   189  				NewMessage([]byte("this is a test message")),
   190  				NewMessage([]byte("and it will work")),
   191  			}
   192  			return m, ackFn, nil
   193  		},
   194  	}
   195  	agi := newAirGapBatchReader(i)
   196  
   197  	err := agi.ConnectWithContext(context.Background())
   198  	assert.NoError(t, err)
   199  
   200  	outMsg, outAckFn, err := agi.ReadWithContext(context.Background())
   201  	assert.NoError(t, err)
   202  	assert.Equal(t, 3, outMsg.Len())
   203  	assert.Equal(t, "hello world", string(outMsg.Get(0).Get()))
   204  	assert.Equal(t, "this is a test message", string(outMsg.Get(1).Get()))
   205  	assert.Equal(t, "and it will work", string(outMsg.Get(2).Get()))
   206  
   207  	assert.NoError(t, outAckFn(context.Background(), response.NewError(errors.New("foobar"))))
   208  	assert.EqualError(t, ackErr, "foobar")
   209  }