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 }