|
| 1 | +package e2e |
| 2 | + |
| 3 | +import ( |
| 4 | + "github.com/stretchr/testify/assert" |
| 5 | + "github.com/twmb/franz-go/pkg/kmsg" |
| 6 | + "sort" |
| 7 | + "testing" |
| 8 | +) |
| 9 | + |
| 10 | +func TestCalculateAppropriateReplicas(t *testing.T) { |
| 11 | + tt := []struct { |
| 12 | + TestName string |
| 13 | + Brokers []kmsg.MetadataResponseBroker |
| 14 | + ReplicationFactor int |
| 15 | + LeaderBroker kmsg.MetadataResponseBroker |
| 16 | + |
| 17 | + // Some cases may have more than one possible solution, each entry in the outer array covers one allowed |
| 18 | + // solution. The compared int32 array order does not matter, except for the very first item as this indicates |
| 19 | + // the preferred leader. For example if you use {2, 0, 1} as expected result this would also be valid for |
| 20 | + // the actual result {2, 1, 0} but not for {1, 2, 0} - because '2' must be the first int32. |
| 21 | + ExpectedResults [][]int32 |
| 22 | + }{ |
| 23 | + { |
| 24 | + TestName: "3 Brokers, no rack, RF = 3", |
| 25 | + Brokers: []kmsg.MetadataResponseBroker{ |
| 26 | + {NodeID: 0, Rack: nil}, |
| 27 | + {NodeID: 1, Rack: nil}, |
| 28 | + {NodeID: 2, Rack: nil}, |
| 29 | + }, |
| 30 | + ReplicationFactor: 3, |
| 31 | + LeaderBroker: kmsg.MetadataResponseBroker{NodeID: 2, Rack: nil}, |
| 32 | + ExpectedResults: [][]int32{{2, 0, 1}}, |
| 33 | + }, |
| 34 | + |
| 35 | + { |
| 36 | + TestName: "3 Brokers, 3 racks, RF = 3", |
| 37 | + Brokers: []kmsg.MetadataResponseBroker{ |
| 38 | + {NodeID: 0, Rack: kmsg.StringPtr("a")}, |
| 39 | + {NodeID: 1, Rack: kmsg.StringPtr("b")}, |
| 40 | + {NodeID: 2, Rack: kmsg.StringPtr("c")}, |
| 41 | + }, |
| 42 | + ReplicationFactor: 3, |
| 43 | + LeaderBroker: kmsg.MetadataResponseBroker{NodeID: 2, Rack: kmsg.StringPtr("c")}, |
| 44 | + ExpectedResults: [][]int32{{2, 0, 1}}, |
| 45 | + }, |
| 46 | + |
| 47 | + { |
| 48 | + TestName: "3 Brokers, 3 racks, RF = 1", |
| 49 | + Brokers: []kmsg.MetadataResponseBroker{ |
| 50 | + {NodeID: 0, Rack: kmsg.StringPtr("a")}, |
| 51 | + {NodeID: 1, Rack: kmsg.StringPtr("b")}, |
| 52 | + {NodeID: 2, Rack: kmsg.StringPtr("c")}, |
| 53 | + }, |
| 54 | + ReplicationFactor: 1, |
| 55 | + LeaderBroker: kmsg.MetadataResponseBroker{NodeID: 1, Rack: kmsg.StringPtr("b")}, |
| 56 | + ExpectedResults: [][]int32{{1}}, |
| 57 | + }, |
| 58 | + |
| 59 | + { |
| 60 | + TestName: "3 Brokers, 3 racks, RF = 2", |
| 61 | + Brokers: []kmsg.MetadataResponseBroker{ |
| 62 | + {NodeID: 0, Rack: kmsg.StringPtr("a")}, |
| 63 | + {NodeID: 1, Rack: kmsg.StringPtr("b")}, |
| 64 | + {NodeID: 2, Rack: kmsg.StringPtr("c")}, |
| 65 | + }, |
| 66 | + ReplicationFactor: 2, |
| 67 | + LeaderBroker: kmsg.MetadataResponseBroker{NodeID: 1, Rack: kmsg.StringPtr("b")}, |
| 68 | + ExpectedResults: [][]int32{{1, 0}, {1, 2}}, |
| 69 | + }, |
| 70 | + |
| 71 | + { |
| 72 | + TestName: "6 Brokers, 3 racks, RF = 3", |
| 73 | + Brokers: []kmsg.MetadataResponseBroker{ |
| 74 | + {NodeID: 0, Rack: kmsg.StringPtr("a")}, |
| 75 | + {NodeID: 1, Rack: kmsg.StringPtr("b")}, |
| 76 | + {NodeID: 2, Rack: kmsg.StringPtr("c")}, |
| 77 | + {NodeID: 3, Rack: kmsg.StringPtr("a")}, |
| 78 | + {NodeID: 4, Rack: kmsg.StringPtr("b")}, |
| 79 | + {NodeID: 5, Rack: kmsg.StringPtr("c")}, |
| 80 | + }, |
| 81 | + ReplicationFactor: 3, |
| 82 | + LeaderBroker: kmsg.MetadataResponseBroker{NodeID: 4, Rack: kmsg.StringPtr("b")}, |
| 83 | + ExpectedResults: [][]int32{{4, 0, 2}, {4, 0, 5}, {4, 3, 2}, {4, 3, 5}}, |
| 84 | + }, |
| 85 | + |
| 86 | + { |
| 87 | + TestName: "4 Brokers, 2 racks, RF = 3", |
| 88 | + Brokers: []kmsg.MetadataResponseBroker{ |
| 89 | + {NodeID: 0, Rack: kmsg.StringPtr("a")}, |
| 90 | + {NodeID: 1, Rack: kmsg.StringPtr("b")}, |
| 91 | + {NodeID: 2, Rack: kmsg.StringPtr("a")}, |
| 92 | + {NodeID: 3, Rack: kmsg.StringPtr("b")}, |
| 93 | + }, |
| 94 | + ReplicationFactor: 3, |
| 95 | + LeaderBroker: kmsg.MetadataResponseBroker{NodeID: 0, Rack: kmsg.StringPtr("a")}, |
| 96 | + ExpectedResults: [][]int32{{0, 1, 2}, {0, 1, 3}, {0, 2, 3}}, |
| 97 | + }, |
| 98 | + |
| 99 | + { |
| 100 | + TestName: "6 Brokers, 3 racks, RF = 3, lowest node id != 0", |
| 101 | + Brokers: []kmsg.MetadataResponseBroker{ |
| 102 | + {NodeID: 10, Rack: kmsg.StringPtr("a")}, |
| 103 | + {NodeID: 11, Rack: kmsg.StringPtr("b")}, |
| 104 | + {NodeID: 12, Rack: kmsg.StringPtr("c")}, |
| 105 | + {NodeID: 13, Rack: kmsg.StringPtr("a")}, |
| 106 | + {NodeID: 14, Rack: kmsg.StringPtr("b")}, |
| 107 | + {NodeID: 15, Rack: kmsg.StringPtr("c")}, |
| 108 | + }, |
| 109 | + ReplicationFactor: 3, |
| 110 | + LeaderBroker: kmsg.MetadataResponseBroker{NodeID: 11, Rack: kmsg.StringPtr("b")}, |
| 111 | + ExpectedResults: [][]int32{{11, 10, 12}, {11, 12, 13}, {11, 13, 15}}, |
| 112 | + }, |
| 113 | + |
| 114 | + { |
| 115 | + TestName: "6 Brokers, 3 racks, RF = 5, lowest node id != 0", |
| 116 | + Brokers: []kmsg.MetadataResponseBroker{ |
| 117 | + {NodeID: 10, Rack: kmsg.StringPtr("a")}, |
| 118 | + {NodeID: 11, Rack: kmsg.StringPtr("b")}, |
| 119 | + {NodeID: 12, Rack: kmsg.StringPtr("c")}, |
| 120 | + {NodeID: 13, Rack: kmsg.StringPtr("a")}, |
| 121 | + {NodeID: 14, Rack: kmsg.StringPtr("b")}, |
| 122 | + {NodeID: 15, Rack: kmsg.StringPtr("c")}, |
| 123 | + }, |
| 124 | + ReplicationFactor: 5, |
| 125 | + LeaderBroker: kmsg.MetadataResponseBroker{NodeID: 11, Rack: kmsg.StringPtr("b")}, |
| 126 | + ExpectedResults: [][]int32{{11, 10, 12, 13, 14}, {11, 10, 13, 14, 15}, {11, 12, 13, 14, 15}, {11, 10, 12, 13, 15}, {11, 10, 12, 14, 15}}, |
| 127 | + }, |
| 128 | + } |
| 129 | + |
| 130 | + svc := Service{} |
| 131 | + for _, test := range tt { |
| 132 | + meta := kmsg.NewMetadataResponse() |
| 133 | + meta.Brokers = test.Brokers |
| 134 | + replicaIDs := svc.calculateAppropriateReplicas(&meta, test.ReplicationFactor, test.LeaderBroker) |
| 135 | + |
| 136 | + matchesAtLeastOneExpectedResult := false |
| 137 | + for _, possibleResult := range test.ExpectedResults { |
| 138 | + isValidResult := possibleResult[0] == replicaIDs[0] && doElementsMatch(possibleResult, replicaIDs) |
| 139 | + if isValidResult { |
| 140 | + matchesAtLeastOneExpectedResult = true |
| 141 | + break |
| 142 | + } |
| 143 | + } |
| 144 | + if !matchesAtLeastOneExpectedResult { |
| 145 | + // Use first elementsmatch to print some valid result along with the actual results. |
| 146 | + assert.ElementsMatch(t, test.ExpectedResults[0], replicaIDs, test.TestName) |
| 147 | + } |
| 148 | + } |
| 149 | +} |
| 150 | + |
| 151 | +func doElementsMatch(a, b []int32) bool { |
| 152 | + if len(a) != len(b) { |
| 153 | + return false |
| 154 | + } |
| 155 | + |
| 156 | + sort.Slice(a, func(i, j int) bool { return a[i] < a[j] }) |
| 157 | + sort.Slice(b, func(i, j int) bool { return b[i] < b[j] }) |
| 158 | + for i, num := range a { |
| 159 | + if num != b[i] { |
| 160 | + return false |
| 161 | + } |
| 162 | + } |
| 163 | + |
| 164 | + return true |
| 165 | +} |
0 commit comments