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

     1  package service
     2  
     3  import (
     4  	"context"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/Jeffail/benthos/v3/lib/log"
     9  	"github.com/Jeffail/benthos/v3/lib/manager"
    10  	"github.com/Jeffail/benthos/v3/lib/metrics"
    11  	"github.com/stretchr/testify/assert"
    12  	"github.com/stretchr/testify/require"
    13  )
    14  
    15  func TestConfigBatching(t *testing.T) {
    16  	spec := NewConfigSpec().
    17  		Field(NewBatchPolicyField("a"))
    18  
    19  	parsedConfig, err := spec.ParseYAML(`
    20  a:
    21    count: 20
    22    period: 5s
    23    processors:
    24      - bloblang: 'root = content().uppercase()'
    25  `, nil)
    26  	require.NoError(t, err)
    27  
    28  	_, err = parsedConfig.FieldTLS("b")
    29  	require.Error(t, err)
    30  
    31  	bConf, err := parsedConfig.FieldBatchPolicy("a")
    32  	require.NoError(t, err)
    33  
    34  	assert.Equal(t, 20, bConf.Count)
    35  	assert.Equal(t, "5s", bConf.Period)
    36  	require.Len(t, bConf.procs, 1)
    37  	assert.Equal(t, "bloblang", bConf.procs[0].Type)
    38  	assert.Equal(t, "root = content().uppercase()", string(bConf.procs[0].Bloblang))
    39  }
    40  
    41  func TestBatcherPeriod(t *testing.T) {
    42  	spec := NewConfigSpec().
    43  		Field(NewBatchPolicyField("a"))
    44  
    45  	parsedConfig, err := spec.ParseYAML(`
    46  a:
    47    period: 300ms
    48    processors:
    49      - bloblang: 'root = content().uppercase()'
    50  `, nil)
    51  	require.NoError(t, err)
    52  
    53  	mgr, err := manager.NewV2(manager.NewResourceConfig(), nil, log.Noop(), metrics.Noop())
    54  	require.NoError(t, err)
    55  
    56  	res := newResourcesFromManager(mgr)
    57  
    58  	bConf, err := parsedConfig.FieldBatchPolicy("a")
    59  	require.NoError(t, err)
    60  
    61  	pol, err := bConf.NewBatcher(res)
    62  	require.NoError(t, err)
    63  
    64  	assert.False(t, pol.Add(NewMessage([]byte("foo"))))
    65  	assert.False(t, pol.Add(NewMessage([]byte("bar"))))
    66  
    67  	v, ok := pol.UntilNext()
    68  	assert.True(t, ok)
    69  	assert.InDelta(t, int(time.Millisecond*300), int(v), float64(time.Millisecond*50))
    70  
    71  	batch, err := pol.Flush(context.Background())
    72  	require.NoError(t, err)
    73  	require.Len(t, batch, 2)
    74  
    75  	bOne, err := batch[0].AsBytes()
    76  	require.NoError(t, err)
    77  	assert.Equal(t, "FOO", string(bOne))
    78  
    79  	bTwo, err := batch[1].AsBytes()
    80  	require.NoError(t, err)
    81  	assert.Equal(t, "BAR", string(bTwo))
    82  
    83  	batch, err = pol.Flush(context.Background())
    84  	require.NoError(t, err)
    85  	require.Len(t, batch, 0)
    86  
    87  	require.NoError(t, pol.Close(context.Background()))
    88  }
    89  
    90  func TestBatcherSize(t *testing.T) {
    91  	spec := NewConfigSpec().
    92  		Field(NewBatchPolicyField("a"))
    93  
    94  	parsedConfig, err := spec.ParseYAML(`
    95  a:
    96    count: 3
    97  `, nil)
    98  	require.NoError(t, err)
    99  
   100  	mgr, err := manager.NewV2(manager.NewResourceConfig(), nil, log.Noop(), metrics.Noop())
   101  	require.NoError(t, err)
   102  
   103  	res := newResourcesFromManager(mgr)
   104  
   105  	bConf, err := parsedConfig.FieldBatchPolicy("a")
   106  	require.NoError(t, err)
   107  
   108  	pol, err := bConf.NewBatcher(res)
   109  	require.NoError(t, err)
   110  
   111  	_, ok := pol.UntilNext()
   112  	assert.False(t, ok)
   113  
   114  	assert.False(t, pol.Add(NewMessage([]byte("foo"))))
   115  	assert.False(t, pol.Add(NewMessage([]byte("bar"))))
   116  	assert.True(t, pol.Add(NewMessage([]byte("baz"))))
   117  
   118  	_, ok = pol.UntilNext()
   119  	assert.False(t, ok)
   120  
   121  	batch, err := pol.Flush(context.Background())
   122  	require.NoError(t, err)
   123  	require.Len(t, batch, 3)
   124  
   125  	bRes, err := batch[0].AsBytes()
   126  	require.NoError(t, err)
   127  	assert.Equal(t, "foo", string(bRes))
   128  
   129  	bRes, err = batch[1].AsBytes()
   130  	require.NoError(t, err)
   131  	assert.Equal(t, "bar", string(bRes))
   132  
   133  	bRes, err = batch[2].AsBytes()
   134  	require.NoError(t, err)
   135  	assert.Equal(t, "baz", string(bRes))
   136  
   137  	batch, err = pol.Flush(context.Background())
   138  	require.NoError(t, err)
   139  	require.Len(t, batch, 0)
   140  
   141  	require.NoError(t, pol.Close(context.Background()))
   142  }