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 //------------------------------------------------------------------------------