@@ -68,6 +68,14 @@ func (i poolTestPeer) freeClientId() string {
68
68
69
69
func (i poolTestPeer ) updateCapacity (uint64 ) {}
70
70
71
+ type poolTestPeerWithCap struct {
72
+ poolTestPeer
73
+
74
+ cap uint64
75
+ }
76
+
77
+ func (i * poolTestPeerWithCap ) updateCapacity (cap uint64 ) { i .cap = cap }
78
+
71
79
func testClientPool (t * testing.T , connLimit , clientCount , paidCount int , randomDisconnect bool ) {
72
80
rand .Seed (time .Now ().UnixNano ())
73
81
var (
@@ -308,9 +316,9 @@ func TestFreeClientKickedOut(t *testing.T) {
308
316
309
317
for i := 0 ; i < 10 ; i ++ {
310
318
pool .connect (poolTestPeer (i ), 1 )
311
- clock .Run (100 * time .Millisecond )
319
+ clock .Run (time .Millisecond )
312
320
}
313
- if pool .connect (poolTestPeer (11 ), 1 ) {
321
+ if pool .connect (poolTestPeer (10 ), 1 ) {
314
322
t .Fatalf ("New free client should be rejected" )
315
323
}
316
324
clock .Run (5 * time .Minute )
@@ -320,8 +328,8 @@ func TestFreeClientKickedOut(t *testing.T) {
320
328
for i := 0 ; i < 10 ; i ++ {
321
329
select {
322
330
case id := <- kicked :
323
- if id != i {
324
- t .Fatalf ("Kicked client mismatch, want %v, got %v" , i , id )
331
+ if id >= 10 {
332
+ t .Fatalf ("Old client should be kicked, now got: %d" , id )
325
333
}
326
334
case <- time .NewTimer (time .Second ).C :
327
335
t .Fatalf ("timeout" )
@@ -364,11 +372,20 @@ func TestDowngradePriorityClient(t *testing.T) {
364
372
pool .setLimits (10 , uint64 (10 )) // Total capacity limit is 10
365
373
pool .setPriceFactors (priceFactors {1 , 0 , 1 }, priceFactors {1 , 0 , 1 })
366
374
367
- pool .addBalance (poolTestPeer (0 ).ID (), uint64 (time .Minute ), false )
368
- pool .connect (poolTestPeer (0 ), 10 )
375
+ p := & poolTestPeerWithCap {
376
+ poolTestPeer : poolTestPeer (0 ),
377
+ }
378
+ pool .addBalance (p .ID (), uint64 (time .Minute ), false )
379
+ pool .connect (p , 10 )
380
+ if p .cap != 10 {
381
+ t .Fatalf ("The capcacity of priority peer hasn't been updated, got: %d" , p .cap )
382
+ }
383
+
369
384
clock .Run (time .Minute ) // All positive balance should be used up.
370
385
time .Sleep (300 * time .Millisecond ) // Ensure the callback is called
371
-
386
+ if p .cap != 1 {
387
+ t .Fatalf ("The capcacity of peer should be downgraded, got: %d" , p .cap )
388
+ }
372
389
pb := pool .ndb .getOrNewPB (poolTestPeer (0 ).ID ())
373
390
if pb .value != 0 {
374
391
t .Fatalf ("Positive balance mismatch, want %v, got %v" , 0 , pb .value )
0 commit comments