github.com/Jeffail/benthos/v3@v3.65.0/lib/broker/round_robin_test.go (about)

     1  package broker
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/Jeffail/benthos/v3/lib/message"
    10  	"github.com/Jeffail/benthos/v3/lib/metrics"
    11  	"github.com/Jeffail/benthos/v3/lib/response"
    12  	"github.com/Jeffail/benthos/v3/lib/types"
    13  )
    14  
    15  var _ types.Consumer = &RoundRobin{}
    16  var _ types.Closable = &RoundRobin{}
    17  
    18  func TestRoundRobinDoubleClose(t *testing.T) {
    19  	oTM, err := NewRoundRobin([]types.Output{}, metrics.Noop())
    20  	if err != nil {
    21  		t.Error(err)
    22  		return
    23  	}
    24  
    25  	// This shouldn't cause a panic
    26  	oTM.CloseAsync()
    27  	oTM.CloseAsync()
    28  }
    29  
    30  //------------------------------------------------------------------------------
    31  
    32  func TestBasicRoundRobin(t *testing.T) {
    33  	nMsgs := 1000
    34  
    35  	outputs := []types.Output{}
    36  	mockOutputs := []*MockOutputType{
    37  		{},
    38  		{},
    39  		{},
    40  	}
    41  
    42  	for _, o := range mockOutputs {
    43  		outputs = append(outputs, o)
    44  	}
    45  
    46  	readChan := make(chan types.Transaction)
    47  	resChan := make(chan types.Response)
    48  
    49  	oTM, err := NewRoundRobin(outputs, metrics.Noop())
    50  	if err != nil {
    51  		t.Error(err)
    52  		return
    53  	}
    54  	if err = oTM.Consume(readChan); err != nil {
    55  		t.Error(err)
    56  		return
    57  	}
    58  
    59  	for i := 0; i < nMsgs; i++ {
    60  		content := [][]byte{[]byte(fmt.Sprintf("hello world %v", i))}
    61  		select {
    62  		case readChan <- types.NewTransaction(message.New(content), resChan):
    63  		case <-time.After(time.Second):
    64  			t.Errorf("Timed out waiting for broker send")
    65  			return
    66  		}
    67  
    68  		go func() {
    69  			var ts types.Transaction
    70  			select {
    71  			case ts = <-mockOutputs[i%3].TChan:
    72  				if !bytes.Equal(ts.Payload.Get(0).Get(), content[0]) {
    73  					t.Errorf("Wrong content returned %s != %s", ts.Payload.Get(0).Get(), content[0])
    74  				}
    75  			case <-mockOutputs[(i+1)%3].TChan:
    76  				t.Errorf("Received message in wrong order: %v != %v", i%3, (i+1)%3)
    77  				return
    78  			case <-mockOutputs[(i+2)%3].TChan:
    79  				t.Errorf("Received message in wrong order: %v != %v", i%3, (i+2)%3)
    80  				return
    81  			case <-time.After(time.Second):
    82  				t.Errorf("Timed out waiting for broker propagate")
    83  				return
    84  			}
    85  
    86  			select {
    87  			case ts.ResponseChan <- response.NewAck():
    88  			case <-time.After(time.Second):
    89  				t.Errorf("Timed out responding to broker")
    90  				return
    91  			}
    92  		}()
    93  
    94  		select {
    95  		case res := <-resChan:
    96  			if res.Error() != nil {
    97  				t.Errorf("Received unexpected errors from broker: %v", res.Error())
    98  			}
    99  		case <-time.After(time.Second):
   100  			t.Errorf("Timed out responding to broker")
   101  			return
   102  		}
   103  	}
   104  
   105  	oTM.CloseAsync()
   106  	if err := oTM.WaitForClose(time.Second * 10); err != nil {
   107  		t.Error(err)
   108  	}
   109  }
   110  
   111  //------------------------------------------------------------------------------
   112  
   113  func BenchmarkBasicRoundRobin(b *testing.B) {
   114  	nOutputs, nMsgs := 3, b.N
   115  
   116  	outputs := []types.Output{}
   117  	mockOutputs := []*MockOutputType{}
   118  
   119  	for i := 0; i < nOutputs; i++ {
   120  		mockOutputs = append(mockOutputs, &MockOutputType{})
   121  		outputs = append(outputs, mockOutputs[i])
   122  	}
   123  
   124  	readChan := make(chan types.Transaction)
   125  	resChan := make(chan types.Response)
   126  
   127  	oTM, err := NewRoundRobin(outputs, metrics.Noop())
   128  	if err != nil {
   129  		b.Error(err)
   130  		return
   131  	}
   132  	if err = oTM.Consume(readChan); err != nil {
   133  		b.Error(err)
   134  		return
   135  	}
   136  
   137  	content := [][]byte{[]byte("hello world")}
   138  
   139  	b.StartTimer()
   140  
   141  	for i := 0; i < nMsgs; i++ {
   142  		readChan <- types.NewTransaction(message.New(content), resChan)
   143  		ts := <-mockOutputs[i%3].TChan
   144  		ts.ResponseChan <- response.NewAck()
   145  		res := <-resChan
   146  		if res.Error() != nil {
   147  			b.Errorf("Received unexpected errors from broker: %v", res.Error())
   148  		}
   149  	}
   150  
   151  	b.StopTimer()
   152  }
   153  
   154  //------------------------------------------------------------------------------