github.com/Jeffail/benthos/v3@v3.65.0/public/service/output_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/types"
    11  	"github.com/stretchr/testify/assert"
    12  )
    13  
    14  type fnOutput struct {
    15  	connect func() error
    16  	write   func(msg *Message) error
    17  	closed  bool
    18  }
    19  
    20  func (f *fnOutput) Connect(ctx context.Context) error {
    21  	return f.connect()
    22  }
    23  
    24  func (f *fnOutput) Write(ctx context.Context, msg *Message) error {
    25  	return f.write(msg)
    26  }
    27  
    28  func (f *fnOutput) Close(ctx context.Context) error {
    29  	f.closed = true
    30  	return nil
    31  }
    32  
    33  func TestOutputAirGapShutdown(t *testing.T) {
    34  	o := &fnOutput{}
    35  	agi := newAirGapWriter(o)
    36  
    37  	err := agi.WaitForClose(time.Millisecond * 5)
    38  	assert.EqualError(t, err, "action timed out")
    39  	assert.False(t, o.closed)
    40  
    41  	agi.CloseAsync()
    42  	err = agi.WaitForClose(time.Millisecond * 5)
    43  	assert.NoError(t, err)
    44  	assert.True(t, o.closed)
    45  }
    46  
    47  func TestOutputAirGapSad(t *testing.T) {
    48  	o := &fnOutput{
    49  		connect: func() error {
    50  			return errors.New("bad connect")
    51  		},
    52  		write: func(m *Message) error {
    53  			return errors.New("bad read")
    54  		},
    55  	}
    56  	agi := newAirGapWriter(o)
    57  
    58  	err := agi.ConnectWithContext(context.Background())
    59  	assert.EqualError(t, err, "bad connect")
    60  
    61  	err = agi.WriteWithContext(context.Background(), message.New(nil))
    62  	assert.EqualError(t, err, "bad read")
    63  
    64  	o.write = func(m *Message) error {
    65  		return ErrNotConnected
    66  	}
    67  
    68  	err = agi.WriteWithContext(context.Background(), message.New(nil))
    69  	assert.Equal(t, types.ErrNotConnected, err)
    70  }
    71  
    72  func TestOutputAirGapHappy(t *testing.T) {
    73  	var wroteMsg string
    74  	o := &fnOutput{
    75  		connect: func() error {
    76  			return nil
    77  		},
    78  		write: func(m *Message) error {
    79  			wroteBytes, _ := m.AsBytes()
    80  			wroteMsg = string(wroteBytes)
    81  			return nil
    82  		},
    83  	}
    84  	agi := newAirGapWriter(o)
    85  
    86  	err := agi.ConnectWithContext(context.Background())
    87  	assert.NoError(t, err)
    88  
    89  	inMsg := message.New([][]byte{[]byte("hello world")})
    90  
    91  	err = agi.WriteWithContext(context.Background(), inMsg)
    92  	assert.NoError(t, err)
    93  
    94  	assert.Equal(t, "hello world", wroteMsg)
    95  }
    96  
    97  type fnBatchOutput struct {
    98  	connect    func() error
    99  	writeBatch func(msgs MessageBatch) error
   100  	closed     bool
   101  }
   102  
   103  func (f *fnBatchOutput) Connect(ctx context.Context) error {
   104  	return f.connect()
   105  }
   106  
   107  func (f *fnBatchOutput) WriteBatch(ctx context.Context, msgs MessageBatch) error {
   108  	return f.writeBatch(msgs)
   109  }
   110  
   111  func (f *fnBatchOutput) Close(ctx context.Context) error {
   112  	f.closed = true
   113  	return nil
   114  }
   115  
   116  func TestBatchOutputAirGapShutdown(t *testing.T) {
   117  	o := &fnBatchOutput{}
   118  	agi := newAirGapBatchWriter(o)
   119  
   120  	err := agi.WaitForClose(time.Millisecond * 5)
   121  	assert.EqualError(t, err, "action timed out")
   122  	assert.False(t, o.closed)
   123  
   124  	agi.CloseAsync()
   125  	err = agi.WaitForClose(time.Millisecond * 5)
   126  	assert.NoError(t, err)
   127  	assert.True(t, o.closed)
   128  }
   129  
   130  func TestBatchOutputAirGapSad(t *testing.T) {
   131  	o := &fnBatchOutput{
   132  		connect: func() error {
   133  			return errors.New("bad connect")
   134  		},
   135  		writeBatch: func(m MessageBatch) error {
   136  			return errors.New("bad read")
   137  		},
   138  	}
   139  	agi := newAirGapBatchWriter(o)
   140  
   141  	err := agi.ConnectWithContext(context.Background())
   142  	assert.EqualError(t, err, "bad connect")
   143  
   144  	err = agi.WriteWithContext(context.Background(), message.New(nil))
   145  	assert.EqualError(t, err, "bad read")
   146  
   147  	o.writeBatch = func(m MessageBatch) error {
   148  		return ErrNotConnected
   149  	}
   150  
   151  	err = agi.WriteWithContext(context.Background(), message.New(nil))
   152  	assert.Equal(t, types.ErrNotConnected, err)
   153  }
   154  
   155  func TestBatchOutputAirGapHappy(t *testing.T) {
   156  	var wroteMsg string
   157  	o := &fnBatchOutput{
   158  		connect: func() error {
   159  			return nil
   160  		},
   161  		writeBatch: func(m MessageBatch) error {
   162  			wroteBytes, _ := m[0].AsBytes()
   163  			wroteMsg = string(wroteBytes)
   164  			return nil
   165  		},
   166  	}
   167  	agi := newAirGapBatchWriter(o)
   168  
   169  	err := agi.ConnectWithContext(context.Background())
   170  	assert.NoError(t, err)
   171  
   172  	inMsg := message.New([][]byte{[]byte("hello world")})
   173  
   174  	err = agi.WriteWithContext(context.Background(), inMsg)
   175  	assert.NoError(t, err)
   176  
   177  	assert.Equal(t, "hello world", wroteMsg)
   178  }