Skip to content

Commit fda37d8

Browse files
asmyasnikovasmyasnikov
committed
fix tests
1 parent 551dbd8 commit fda37d8

File tree

2 files changed

+49
-30
lines changed

2 files changed

+49
-30
lines changed

‎internal/balancer/connections_state.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,9 @@ func (s *connectionsState) selectRandomConnection(conns []conn.Conn, allowBanned
119119
}
120120

121121
func connsToNodeIDMap(conns []conn.Conn) (nodes map[uint32]conn.Conn) {
122+
if len(conns) == 0 {
123+
return nil
124+
}
122125
nodes = make(map[uint32]conn.Conn, len(conns))
123126
for _, c := range conns {
124127
nodes[c.Endpoint().NodeID()] = c

‎internal/balancer/connections_state_test.go‎

Lines changed: 46 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,9 @@ func TestConnsToNodeIDMap(t *testing.T) {
2828
source: []conn.Conn{
2929
&mock.Conn{NodeIDField: 0},
3030
},
31-
res: nil,
31+
res: map[uint32]conn.Conn{
32+
0: &mock.Conn{NodeIDField: 0},
33+
},
3234
},
3335
{
3436
name: "NonZero",
@@ -49,6 +51,7 @@ func TestConnsToNodeIDMap(t *testing.T) {
4951
&mock.Conn{NodeIDField: 10},
5052
},
5153
res: map[uint32]conn.Conn{
54+
0: &mock.Conn{NodeIDField: 0},
5255
1: &mock.Conn{NodeIDField: 1},
5356
10: &mock.Conn{NodeIDField: 10},
5457
},
@@ -247,70 +250,83 @@ func TestNewState(t *testing.T) {
247250
{
248251
name: "NoFilter",
249252
state: newConnectionsState([]conn.Conn{
250-
&mock.Conn{AddrField: "1"},
251-
&mock.Conn{AddrField: "2"},
253+
&mock.Conn{AddrField: "1", NodeIDField: 1},
254+
&mock.Conn{AddrField: "2", NodeIDField: 2},
252255
}, nil, balancerConfig.Info{}, false),
253256
res: &connectionsState{
254-
connByNodeID: nil,
257+
connByNodeID: map[uint32]conn.Conn{
258+
1: &mock.Conn{AddrField: "1", NodeIDField: 1},
259+
2: &mock.Conn{AddrField: "2", NodeIDField: 2},
260+
},
255261
prefer: []conn.Conn{
256-
&mock.Conn{AddrField: "1"},
257-
&mock.Conn{AddrField: "2"},
262+
&mock.Conn{AddrField: "1", NodeIDField: 1},
263+
&mock.Conn{AddrField: "2", NodeIDField: 2},
258264
},
259265
fallback: nil,
260266
all: []conn.Conn{
261-
&mock.Conn{AddrField: "1"},
262-
&mock.Conn{AddrField: "2"},
267+
&mock.Conn{AddrField: "1", NodeIDField: 1},
268+
&mock.Conn{AddrField: "2", NodeIDField: 2},
263269
},
264270
},
265271
},
266272
{
267273
name: "FilterDenyFallback",
268274
state: newConnectionsState([]conn.Conn{
269-
&mock.Conn{AddrField: "t1", LocationField: "t"},
270-
&mock.Conn{AddrField: "f1", LocationField: "f"},
271-
&mock.Conn{AddrField: "t2", LocationField: "t"},
272-
&mock.Conn{AddrField: "f2", LocationField: "f"},
275+
&mock.Conn{AddrField: "t1", NodeIDField: 1, LocationField: "t"},
276+
&mock.Conn{AddrField: "f1", NodeIDField: 2, LocationField: "f"},
277+
&mock.Conn{AddrField: "t2", NodeIDField: 3, LocationField: "t"},
278+
&mock.Conn{AddrField: "f2", NodeIDField: 4, LocationField: "f"},
273279
}, func(info balancerConfig.Info, c conn.Conn) bool {
274280
return info.SelfLocation == c.Endpoint().Location()
275281
}, balancerConfig.Info{SelfLocation: "t"}, false),
276282
res: &connectionsState{
277-
connByNodeID: nil,
283+
connByNodeID: map[uint32]conn.Conn{
284+
1: &mock.Conn{AddrField: "t1", NodeIDField: 1, LocationField: "t"},
285+
2: &mock.Conn{AddrField: "f1", NodeIDField: 2, LocationField: "f"},
286+
3: &mock.Conn{AddrField: "t2", NodeIDField: 3, LocationField: "t"},
287+
4: &mock.Conn{AddrField: "f2", NodeIDField: 4, LocationField: "f"},
288+
},
278289
prefer: []conn.Conn{
279-
&mock.Conn{AddrField: "t1", LocationField: "t"},
280-
&mock.Conn{AddrField: "t2", LocationField: "t"},
290+
&mock.Conn{AddrField: "t1", NodeIDField: 1, LocationField: "t"},
291+
&mock.Conn{AddrField: "t2", NodeIDField: 3, LocationField: "t"},
281292
},
282293
fallback: nil,
283294
all: []conn.Conn{
284-
&mock.Conn{AddrField: "t1", LocationField: "t"},
285-
&mock.Conn{AddrField: "t2", LocationField: "t"},
295+
&mock.Conn{AddrField: "t1", NodeIDField: 1, LocationField: "t"},
296+
&mock.Conn{AddrField: "t2", NodeIDField: 3, LocationField: "t"},
286297
},
287298
},
288299
},
289300
{
290301
name: "FilterAllowFallback",
291302
state: newConnectionsState([]conn.Conn{
292-
&mock.Conn{AddrField: "t1", LocationField: "t"},
293-
&mock.Conn{AddrField: "f1", LocationField: "f"},
294-
&mock.Conn{AddrField: "t2", LocationField: "t"},
295-
&mock.Conn{AddrField: "f2", LocationField: "f"},
303+
&mock.Conn{AddrField: "t1", NodeIDField: 1, LocationField: "t"},
304+
&mock.Conn{AddrField: "f1", NodeIDField: 2, LocationField: "f"},
305+
&mock.Conn{AddrField: "t2", NodeIDField: 3, LocationField: "t"},
306+
&mock.Conn{AddrField: "f2", NodeIDField: 4, LocationField: "f"},
296307
}, func(info balancerConfig.Info, c conn.Conn) bool {
297308
return info.SelfLocation == c.Endpoint().Location()
298309
}, balancerConfig.Info{SelfLocation: "t"}, true),
299310
res: &connectionsState{
300-
connByNodeID: nil,
311+
connByNodeID: map[uint32]conn.Conn{
312+
1: &mock.Conn{AddrField: "t1", NodeIDField: 1, LocationField: "t"},
313+
2: &mock.Conn{AddrField: "f1", NodeIDField: 2, LocationField: "f"},
314+
3: &mock.Conn{AddrField: "t2", NodeIDField: 3, LocationField: "t"},
315+
4: &mock.Conn{AddrField: "f2", NodeIDField: 4, LocationField: "f"},
316+
},
301317
prefer: []conn.Conn{
302-
&mock.Conn{AddrField: "t1", LocationField: "t"},
303-
&mock.Conn{AddrField: "t2", LocationField: "t"},
318+
&mock.Conn{AddrField: "t1", NodeIDField: 1, LocationField: "t"},
319+
&mock.Conn{AddrField: "t2", NodeIDField: 3, LocationField: "t"},
304320
},
305321
fallback: []conn.Conn{
306-
&mock.Conn{AddrField: "f1", LocationField: "f"},
307-
&mock.Conn{AddrField: "f2", LocationField: "f"},
322+
&mock.Conn{AddrField: "f1", NodeIDField: 2, LocationField: "f"},
323+
&mock.Conn{AddrField: "f2", NodeIDField: 4, LocationField: "f"},
308324
},
309325
all: []conn.Conn{
310-
&mock.Conn{AddrField: "t1", LocationField: "t"},
311-
&mock.Conn{AddrField: "f1", LocationField: "f"},
312-
&mock.Conn{AddrField: "t2", LocationField: "t"},
313-
&mock.Conn{AddrField: "f2", LocationField: "f"},
326+
&mock.Conn{AddrField: "t1", NodeIDField: 1, LocationField: "t"},
327+
&mock.Conn{AddrField: "f1", NodeIDField: 2, LocationField: "f"},
328+
&mock.Conn{AddrField: "t2", NodeIDField: 3, LocationField: "t"},
329+
&mock.Conn{AddrField: "f2", NodeIDField: 4, LocationField: "f"},
314330
},
315331
},
316332
},

0 commit comments

Comments
 (0)