github.com/Jeffail/benthos/v3@v3.65.0/lib/message/batch/policy_test.go (about)

     1  package batch
     2  
     3  import (
     4  	"reflect"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/Jeffail/benthos/v3/lib/condition"
     9  	"github.com/Jeffail/benthos/v3/lib/log"
    10  	"github.com/Jeffail/benthos/v3/lib/message"
    11  	"github.com/Jeffail/benthos/v3/lib/metrics"
    12  	"github.com/Jeffail/benthos/v3/lib/processor"
    13  	"github.com/stretchr/testify/assert"
    14  	"github.com/stretchr/testify/require"
    15  )
    16  
    17  func TestPolicyNoop(t *testing.T) {
    18  	conf := NewPolicyConfig()
    19  	assert.True(t, conf.IsNoop())
    20  
    21  	conf = NewPolicyConfig()
    22  	conf.Count = 2
    23  	assert.False(t, conf.IsNoop())
    24  
    25  	conf = NewPolicyConfig()
    26  	conf.Condition = condition.NewConfig()
    27  	assert.False(t, conf.IsNoop())
    28  
    29  	conf = NewPolicyConfig()
    30  	conf.Check = "foo.bar"
    31  	assert.False(t, conf.IsNoop())
    32  
    33  	conf = NewPolicyConfig()
    34  	conf.ByteSize = 10
    35  	assert.False(t, conf.IsNoop())
    36  
    37  	conf = NewPolicyConfig()
    38  	conf.Period = "10s"
    39  	assert.False(t, conf.IsNoop())
    40  }
    41  
    42  func TestPolicyBasic(t *testing.T) {
    43  	conf := NewPolicyConfig()
    44  	conf.Count = 2
    45  	conf.ByteSize = 0
    46  
    47  	pol, err := NewPolicy(conf, nil, log.Noop(), metrics.Noop())
    48  	require.NoError(t, err)
    49  
    50  	t.Cleanup(func() {
    51  		pol.CloseAsync()
    52  		require.NoError(t, pol.WaitForClose(time.Second))
    53  	})
    54  
    55  	if v := pol.UntilNext(); v >= 0 {
    56  		t.Errorf("Non-negative period: %v", v)
    57  	}
    58  
    59  	if exp, act := 0, pol.Count(); exp != act {
    60  		t.Errorf("Wrong count: %v != %v", act, exp)
    61  	}
    62  
    63  	exp := [][]byte{[]byte("foo"), []byte("bar")}
    64  
    65  	if pol.Add(message.NewPart(exp[0])) {
    66  		t.Error("Unexpected batch")
    67  	}
    68  	if exp, act := 1, pol.Count(); exp != act {
    69  		t.Errorf("Wrong count: %v != %v", act, exp)
    70  	}
    71  	if !pol.Add(message.NewPart(exp[1])) {
    72  		t.Error("Expected batch")
    73  	}
    74  	if exp, act := 2, pol.Count(); exp != act {
    75  		t.Errorf("Wrong count: %v != %v", act, exp)
    76  	}
    77  
    78  	msg := pol.Flush()
    79  	if !reflect.DeepEqual(exp, message.GetAllBytes(msg)) {
    80  		t.Errorf("Wrong result: %s != %s", message.GetAllBytes(msg), exp)
    81  	}
    82  	if exp, act := 0, pol.Count(); exp != act {
    83  		t.Errorf("Wrong count: %v != %v", act, exp)
    84  	}
    85  
    86  	if msg = pol.Flush(); msg != nil {
    87  		t.Error("Non-nil empty flush")
    88  	}
    89  }
    90  
    91  func TestPolicyPeriod(t *testing.T) {
    92  	conf := NewPolicyConfig()
    93  	conf.Period = "300ms"
    94  
    95  	pol, err := NewPolicy(conf, nil, log.Noop(), metrics.Noop())
    96  	if err != nil {
    97  		t.Fatal(err)
    98  	}
    99  
   100  	t.Cleanup(func() {
   101  		pol.CloseAsync()
   102  		require.NoError(t, pol.WaitForClose(time.Second))
   103  	})
   104  
   105  	if pol.Add(message.NewPart(nil)) {
   106  		t.Error("Unexpected batch ready")
   107  	}
   108  
   109  	if v := pol.UntilNext(); v >= (time.Millisecond*300) || v < (time.Millisecond*100) {
   110  		t.Errorf("Wrong period: %v", v)
   111  	}
   112  
   113  	<-time.After(time.Millisecond * 500)
   114  	if v := pol.UntilNext(); v >= (time.Millisecond * 100) {
   115  		t.Errorf("Wrong period: %v", v)
   116  	}
   117  
   118  	if v := pol.Flush(); v == nil {
   119  		t.Error("Nil msgs from flush")
   120  	}
   121  
   122  	if v := pol.UntilNext(); v >= (time.Millisecond*300) || v < (time.Millisecond*100) {
   123  		t.Errorf("Wrong period: %v", v)
   124  	}
   125  }
   126  
   127  func TestPolicySize(t *testing.T) {
   128  	conf := NewPolicyConfig()
   129  	conf.ByteSize = 10
   130  
   131  	pol, err := NewPolicy(conf, nil, log.Noop(), metrics.Noop())
   132  	if err != nil {
   133  		t.Fatal(err)
   134  	}
   135  
   136  	t.Cleanup(func() {
   137  		pol.CloseAsync()
   138  		require.NoError(t, pol.WaitForClose(time.Second))
   139  	})
   140  
   141  	exp := [][]byte{[]byte("foo bar"), []byte("baz qux")}
   142  
   143  	if pol.Add(message.NewPart(exp[0])) {
   144  		t.Error("Unexpected batch")
   145  	}
   146  	if !pol.Add(message.NewPart(exp[1])) {
   147  		t.Error("Expected batch")
   148  	}
   149  
   150  	msg := pol.Flush()
   151  	if !reflect.DeepEqual(exp, message.GetAllBytes(msg)) {
   152  		t.Errorf("Wrong result: %s != %s", message.GetAllBytes(msg), exp)
   153  	}
   154  
   155  	if msg = pol.Flush(); msg != nil {
   156  		t.Error("Non-nil empty flush")
   157  	}
   158  }
   159  
   160  func TestPolicyCheck(t *testing.T) {
   161  	conf := NewPolicyConfig()
   162  	conf.Check = `content() == "bar"`
   163  
   164  	pol, err := NewPolicy(conf, nil, log.Noop(), metrics.Noop())
   165  	if err != nil {
   166  		t.Fatal(err)
   167  	}
   168  
   169  	t.Cleanup(func() {
   170  		pol.CloseAsync()
   171  		require.NoError(t, pol.WaitForClose(time.Second))
   172  	})
   173  
   174  	exp := [][]byte{[]byte("foo"), []byte("bar")}
   175  
   176  	if pol.Add(message.NewPart(exp[0])) {
   177  		t.Error("Unexpected batch")
   178  	}
   179  	if !pol.Add(message.NewPart(exp[1])) {
   180  		t.Error("Expected batch")
   181  	}
   182  
   183  	msg := pol.Flush()
   184  	if !reflect.DeepEqual(exp, message.GetAllBytes(msg)) {
   185  		t.Errorf("Wrong result: %s != %s", message.GetAllBytes(msg), exp)
   186  	}
   187  
   188  	if msg = pol.Flush(); msg != nil {
   189  		t.Error("Non-nil empty flush")
   190  	}
   191  }
   192  
   193  func TestPolicyCheckAdvanced(t *testing.T) {
   194  	conf := NewPolicyConfig()
   195  	conf.Check = `batch_size() >= 3`
   196  
   197  	pol, err := NewPolicy(conf, nil, log.Noop(), metrics.Noop())
   198  	if err != nil {
   199  		t.Fatal(err)
   200  	}
   201  
   202  	t.Cleanup(func() {
   203  		pol.CloseAsync()
   204  		require.NoError(t, pol.WaitForClose(time.Second))
   205  	})
   206  
   207  	exp := [][]byte{[]byte("foo"), []byte("bar"), []byte("baz")}
   208  
   209  	if pol.Add(message.NewPart(exp[0])) {
   210  		t.Error("Unexpected batch")
   211  	}
   212  	if pol.Add(message.NewPart(exp[1])) {
   213  		t.Error("Expected batch")
   214  	}
   215  	if !pol.Add(message.NewPart(exp[2])) {
   216  		t.Error("Expected batch")
   217  	}
   218  
   219  	msg := pol.Flush()
   220  	if !reflect.DeepEqual(exp, message.GetAllBytes(msg)) {
   221  		t.Errorf("Wrong result: %s != %s", message.GetAllBytes(msg), exp)
   222  	}
   223  
   224  	if msg = pol.Flush(); msg != nil {
   225  		t.Error("Non-nil empty flush")
   226  	}
   227  }
   228  
   229  func TestPolicyCondition(t *testing.T) {
   230  	cond := condition.NewConfig()
   231  	cond.Type = condition.TypeText
   232  	cond.Text.Operator = "equals"
   233  	cond.Text.Arg = "bar"
   234  
   235  	conf := NewPolicyConfig()
   236  	conf.Condition = cond
   237  
   238  	pol, err := NewPolicy(conf, nil, log.Noop(), metrics.Noop())
   239  	if err != nil {
   240  		t.Fatal(err)
   241  	}
   242  
   243  	t.Cleanup(func() {
   244  		pol.CloseAsync()
   245  		require.NoError(t, pol.WaitForClose(time.Second))
   246  	})
   247  
   248  	exp := [][]byte{[]byte("foo"), []byte("bar")}
   249  
   250  	if pol.Add(message.NewPart(exp[0])) {
   251  		t.Error("Unexpected batch")
   252  	}
   253  	if !pol.Add(message.NewPart(exp[1])) {
   254  		t.Error("Expected batch")
   255  	}
   256  
   257  	msg := pol.Flush()
   258  	if !reflect.DeepEqual(exp, message.GetAllBytes(msg)) {
   259  		t.Errorf("Wrong result: %s != %s", message.GetAllBytes(msg), exp)
   260  	}
   261  
   262  	if msg = pol.Flush(); msg != nil {
   263  		t.Error("Non-nil empty flush")
   264  	}
   265  }
   266  
   267  func TestPolicyArchived(t *testing.T) {
   268  	conf := NewPolicyConfig()
   269  	conf.Count = 2
   270  	conf.ByteSize = 0
   271  
   272  	procConf := processor.NewConfig()
   273  	procConf.Type = processor.TypeArchive
   274  	procConf.Archive.Format = "lines"
   275  
   276  	conf.Processors = append(conf.Processors, procConf)
   277  
   278  	pol, err := NewPolicy(conf, nil, log.Noop(), metrics.Noop())
   279  	require.NoError(t, err)
   280  
   281  	t.Cleanup(func() {
   282  		pol.CloseAsync()
   283  		require.NoError(t, pol.WaitForClose(time.Second))
   284  	})
   285  
   286  	exp := [][]byte{[]byte("foo\nbar")}
   287  
   288  	assert.False(t, pol.Add(message.NewPart([]byte("foo"))))
   289  	assert.Equal(t, 1, pol.Count())
   290  
   291  	assert.True(t, pol.Add(message.NewPart([]byte("bar"))))
   292  	assert.Equal(t, 2, pol.Count())
   293  
   294  	msg := pol.Flush()
   295  	assert.Equal(t, exp, message.GetAllBytes(msg))
   296  	assert.Equal(t, 0, pol.Count())
   297  
   298  	msg = pol.Flush()
   299  	assert.Nil(t, msg)
   300  }
   301  
   302  func TestPolicySplit(t *testing.T) {
   303  	conf := NewPolicyConfig()
   304  	conf.Count = 2
   305  	conf.ByteSize = 0
   306  
   307  	procConf := processor.NewConfig()
   308  	procConf.Type = processor.TypeSplit
   309  
   310  	conf.Processors = append(conf.Processors, procConf)
   311  
   312  	pol, err := NewPolicy(conf, nil, log.Noop(), metrics.Noop())
   313  	if err != nil {
   314  		t.Fatal(err)
   315  	}
   316  
   317  	t.Cleanup(func() {
   318  		pol.CloseAsync()
   319  		require.NoError(t, pol.WaitForClose(time.Second))
   320  	})
   321  
   322  	exp := [][]byte{[]byte("foo"), []byte("bar")}
   323  
   324  	if pol.Add(message.NewPart([]byte("foo"))) {
   325  		t.Error("Unexpected batch")
   326  	}
   327  	if exp, act := 1, pol.Count(); exp != act {
   328  		t.Errorf("Wrong count: %v != %v", act, exp)
   329  	}
   330  	if !pol.Add(message.NewPart([]byte("bar"))) {
   331  		t.Error("Expected batch")
   332  	}
   333  	if exp, act := 2, pol.Count(); exp != act {
   334  		t.Errorf("Wrong count: %v != %v", act, exp)
   335  	}
   336  
   337  	msg := pol.Flush()
   338  	if !reflect.DeepEqual(exp, message.GetAllBytes(msg)) {
   339  		t.Errorf("Wrong result: %s != %s", message.GetAllBytes(msg), exp)
   340  	}
   341  	if exp, act := 0, pol.Count(); exp != act {
   342  		t.Errorf("Wrong count: %v != %v", act, exp)
   343  	}
   344  
   345  	if msg = pol.Flush(); msg != nil {
   346  		t.Error("Non-nil empty flush")
   347  	}
   348  }
   349  
   350  func TestPolicySplitAny(t *testing.T) {
   351  	conf := NewPolicyConfig()
   352  	conf.Count = 2
   353  	conf.ByteSize = 0
   354  
   355  	procConf := processor.NewConfig()
   356  	procConf.Type = processor.TypeSplit
   357  
   358  	conf.Processors = append(conf.Processors, procConf)
   359  
   360  	pol, err := NewPolicy(conf, nil, log.Noop(), metrics.Noop())
   361  	if err != nil {
   362  		t.Fatal(err)
   363  	}
   364  
   365  	t.Cleanup(func() {
   366  		pol.CloseAsync()
   367  		require.NoError(t, pol.WaitForClose(time.Second))
   368  	})
   369  
   370  	if pol.Add(message.NewPart([]byte("foo"))) {
   371  		t.Error("Unexpected batch")
   372  	}
   373  	if exp, act := 1, pol.Count(); exp != act {
   374  		t.Errorf("Wrong count: %v != %v", act, exp)
   375  	}
   376  	if !pol.Add(message.NewPart([]byte("bar"))) {
   377  		t.Error("Expected batch")
   378  	}
   379  	if exp, act := 2, pol.Count(); exp != act {
   380  		t.Errorf("Wrong count: %v != %v", act, exp)
   381  	}
   382  
   383  	msgs := pol.FlushAny()
   384  	assert.Equal(t, 2, len(msgs))
   385  	assert.Equal(t, 1, msgs[0].Len())
   386  	assert.Equal(t, 1, msgs[1].Len())
   387  	assert.Equal(t, "foo", string(msgs[0].Get(0).Get()))
   388  	assert.Equal(t, "bar", string(msgs[1].Get(0).Get()))
   389  }
   390  
   391  func TestPolicySplitAnyDeprecated(t *testing.T) {
   392  	conf := NewPolicyConfig()
   393  	conf.Count = 2
   394  	conf.ByteSize = 0
   395  
   396  	procConf := processor.NewConfig()
   397  	procConf.Type = processor.TypeSplit
   398  
   399  	conf.Processors = append(conf.Processors, procConf)
   400  
   401  	pol, err := NewPolicy(conf, nil, log.Noop(), metrics.Noop())
   402  	if err != nil {
   403  		t.Fatal(err)
   404  	}
   405  
   406  	t.Cleanup(func() {
   407  		pol.CloseAsync()
   408  		require.NoError(t, pol.WaitForClose(time.Second))
   409  	})
   410  
   411  	if pol.Add(message.NewPart([]byte("foo"))) {
   412  		t.Error("Unexpected batch")
   413  	}
   414  	if exp, act := 1, pol.Count(); exp != act {
   415  		t.Errorf("Wrong count: %v != %v", act, exp)
   416  	}
   417  	if !pol.Add(message.NewPart([]byte("bar"))) {
   418  		t.Error("Expected batch")
   419  	}
   420  	if exp, act := 2, pol.Count(); exp != act {
   421  		t.Errorf("Wrong count: %v != %v", act, exp)
   422  	}
   423  
   424  	msgs := pol.FlushAny()
   425  	assert.Equal(t, 2, len(msgs))
   426  	assert.Equal(t, 1, msgs[0].Len())
   427  	assert.Equal(t, 1, msgs[1].Len())
   428  	assert.Equal(t, "foo", string(msgs[0].Get(0).Get()))
   429  	assert.Equal(t, "bar", string(msgs[1].Get(0).Get()))
   430  }